ruspa/dynamic_tree.go

package main

import (
	"errors"
	"fmt"
	"path/filepath"
	"strings"
)

const (
	TopType          = "top"
	IntermediateType = "intermediate"
	BottomType       = "bottom"
	SingleType       = "single"
)

type ByteSize uint

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
	GB
	TB
	PB
)

var ErrUnknownType = errors.New("unknown type")
var ErrShouldBeBottom = errors.New("should be of type Bottom")
var ErrChannelError = errors.New("channel errored")
var ErrUnknownUnit = errors.New("unknown unit")
var (
	NodeB   = "│"
	NodeT   = "├──"
	NodeL   = "└──"
	NodePad = "   "
)

type Node interface {
	SetUnit(string)
	GetUnit() ByteSize
	Size() int64
	Spawn(int) error
	Collect() error
	Collector() chan int64
	Level() int
}

func isLastLevel(node Node, maxLevel int) bool {
	return maxLevel-1 == node.Level()
}

func setUnit(unit string) ByteSize {
	switch unit {
	case "B":
		return B
	case "KB":
		return KB
	case "MB":
		return MB
	case "GB":
		return GB
	case "TB":
		return TB
	case "PB":
		// I reckon this will never happen, but you never know
		return PB
	default:
		panic(ErrUnknownUnit)
	}
}

func fmtUnit(unit ByteSize) string {
	switch unit {
	case B:
		return "B"
	case KB:
		return "KB"
	case MB:
		return "MB"
	case GB:
		return "GB"
	case TB:
		return "TB"
	case PB:
		// I reckon this will never happen, but you never know
		return "PB"
	default:
		panic(ErrUnknownUnit)
	}
}

func fmtSize(size int64, unit ByteSize) string {
	var dimension float64
	dimension = float64(size) / float64(unit)
	return fmt.Sprintf("%.2f %s", dimension, fmtUnit(unit))
}

//////////////////////////////
/////////// Top //////////////
//////////////////////////////

type Top struct {
	path    string
	size    int64
	unit    ByteSize
	tree    []Node
	collect chan int64
	level   int
}

func NewTop(path string) *Top {
	t := &Top{path: path, level: 0, unit: KB}
	t.collect = make(chan int64)
	return t
}

func (t *Top) SetUnit(unit string) {
	t.unit = setUnit(unit)
}

func (t *Top) GetUnit() ByteSize {
	return t.unit
}

func (t *Top) Size() int64 {
	return t.size
}

func (t *Top) Spawn(maxLevel int) error {
	if t.Level() >= maxLevel {
		return ErrShouldBeBottom
	}
	files, err := ListDir(t.path)
	if err != nil {
		return err
	}

	for _, info := range files {
		switch mode := info.Mode(); {
		case mode.IsDir():
			var child Node
			path := filepath.Join(t.path, info.Name())
			if isLastLevel(t, maxLevel) || !AnyDirectoryDownThere(path) {
				child = t.newBottom(path)
			} else {
				child = t.newIntermediate(path)
			}
			go child.Spawn(maxLevel)
			defer func() {
				go child.Collect()
			}()
		case !mode.IsDir() && mode.IsRegular():
			t.collect <- info.Size()
		}
	}
	return nil
}

func (t *Top) Collect() error {
	for s := range t.collect {
		t.size += s
	}
	return nil
}

func (t *Top) Collector() chan int64 {
	return t.collect
}

func (t *Top) Level() int {
	return t.level
}

func (t *Top) Name() string {
	return filepath.Base(t.path)
}

func (t *Top) String() string {
	var out string
	var lines []string
	out += fmt.Sprintf("(%s) %s\n", fmtSize(t.Size(), t.GetUnit()), t.path)
	treeSize := len(t.tree)
	for _, child := range t.tree {
		lines = append(lines, fmt.Sprintf("%s", child))
	}
	for a := 0; a < treeSize-1; a++ {
		childLines := strings.Split(lines[a], "\n")
		out += fmt.Sprintf(" %s%s\n", NodeT, childLines[0])
		for b := 1; b < len(childLines); b++ {
			out += fmt.Sprintf(" %s%s\n", NodeB, childLines[b])
		}
	}
	childLines := strings.Split(lines[treeSize-1], "\n")
	out += fmt.Sprintf(" %s%s\n", NodeL, childLines[0])
	for a := 1; a < len(childLines); a++ {
		out += fmt.Sprintf("%s%s\n", "  ", childLines[a])
	}
	return out
}

func (t *Top) addToTree(child Node) {
	t.tree = append(t.tree, child)
}

func (t *Top) newChild(path, kind string) Node {
	switch {
	case kind == IntermediateType:
		i := NewIntermediate(path, t)
		t.addToTree(i)
		return i
	case kind == BottomType:
		b := NewBottom(path, t)
		t.addToTree(b)
		return b
	default:
		panic(ErrUnknownType)
	}
}

func (t *Top) newIntermediate(path string) *Intermediate {
	return t.newChild(path, IntermediateType).(*Intermediate)
}

func (t *Top) newBottom(path string) *Bottom {
	return t.newChild(path, BottomType).(*Bottom)
}

//////////////////////////////
/////// Intermediate /////////
//////////////////////////////

type Intermediate struct {
	path    string
	size    int64
	unit    ByteSize
	parent  Node
	tree    []Node
	collect chan int64
	refer   chan int64
	level   int
}

func NewIntermediate(path string, parent Node) *Intermediate {
	i := &Intermediate{path: path, refer: parent.Collector(), parent: parent}
	i.collect = make(chan int64)
	i.level = parent.Level() + 1
	i.unit = parent.GetUnit()
	return i
}

func (i *Intermediate) SetUnit(unit string) {
	i.unit = setUnit(unit)
}

func (i *Intermediate) GetUnit() ByteSize {
	return i.unit
}

func (i *Intermediate) Size() int64 {
	return i.size
}

func (i *Intermediate) Spawn(maxLevel int) error {
	if i.Level() >= maxLevel {
		return ErrShouldBeBottom
	}

	files, err := ListDir(i.path)
	if err != nil {
		return err
	}

	for _, info := range files {
		switch mode := info.Mode(); {
		case mode.IsDir():
			var child Node
			path := filepath.Join(i.path, info.Name())
			if isLastLevel(i, maxLevel) || !AnyDirectoryDownThere(path) {
				child = i.newBottom(path)
			} else {
				child = i.newIntermediate(path)
			}
			go child.Spawn(maxLevel)
			defer func() {
				go child.Collect()
			}()
		case !mode.IsDir() && mode.IsRegular():
			i.collect <- info.Size()
		}
	}

	return nil
}

func (i *Intermediate) Collect() error {
	for s := range i.collect {
		i.size += s
		i.refer <- s
	}
	close(i.refer)
	return nil
}

func (i *Intermediate) Collector() chan int64 {
	return i.collect
}

func (i *Intermediate) Level() int {
	return i.level
}

func (i *Intermediate) Name() string {
	return filepath.Base(i.path)
}

func (i *Intermediate) String() string {
	var lines []string
	out := fmt.Sprintf("(%s) %s\n", fmtSize(i.Size(), i.GetUnit()), i.Name())
	treeSize := len(i.tree)
	for _, child := range i.tree {
		lines = append(lines, fmt.Sprintf("%s", child))
	}

	for a := 0; a < treeSize-1; a++ {
		childLines := strings.Split(lines[a], "\n")
		out += fmt.Sprintf("%s%s%s\n", NodePad, NodeT, childLines[0])
		for b := 1; b < len(childLines); b++ {
			out += fmt.Sprintf("%s%s%s\n", NodePad, NodeB, childLines[b])
		}
	}

	childLines := strings.Split(lines[treeSize-1], "\n")
	lenChildLines := len(childLines)
	if lenChildLines > 1 {
		out += fmt.Sprintf("%s%s%s\n", NodePad, NodeL, childLines[0])
		for a := 1; a < lenChildLines-2; a++ {
			out += fmt.Sprintf(" %s%s\n", NodePad, childLines[a])
		}
		out += fmt.Sprintf(" %s%s", NodePad, childLines[lenChildLines-1])
	} else {
		out += fmt.Sprintf("%s%s%s", NodePad, NodeL, childLines[0])
	}
	return out
}

func (i *Intermediate) addToTree(child Node) {
	i.tree = append(i.tree, child)
}

func (i *Intermediate) newChild(path, kind string) Node {
	switch {
	case kind == IntermediateType:
		c := NewIntermediate(path, i)
		i.addToTree(c)
		return c
	case kind == BottomType:
		b := NewBottom(path, i)
		i.addToTree(b)
		return b
	default:
		panic(ErrUnknownType)
	}
}

func (i *Intermediate) newIntermediate(path string) *Intermediate {
	return i.newChild(path, IntermediateType).(*Intermediate)
}

func (i *Intermediate) newBottom(path string) *Bottom {
	return i.newChild(path, BottomType).(*Bottom)
}

//////////////////////////////
////////// Bottom ////////////
//////////////////////////////

type Bottom struct {
	path    string
	size    int64
	unit    ByteSize
	parent  Node
	walker  Walker
	collect chan int64
	refer   chan int64
	level   int
}

func NewBottom(path string, parent Node) *Bottom {
	b := &Bottom{path: path, refer: parent.Collector(), parent: parent, unit: parent.GetUnit()}
	b.collect = make(chan int64)
	b.level = parent.Level() + 1
	b.unit = parent.GetUnit()
	return b
}

func (b *Bottom) SetUnit(unit string) {
	b.unit = setUnit(unit)
}

func (b *Bottom) GetUnit() ByteSize {
	return b.unit
}

func (b *Bottom) Size() int64 {
	return b.size
}

func (b *Bottom) Spawn(maxLevel int) error {
	b.walker = NewRealWalker(b.path, b.collect)
	go b.walker.Walk()
	return nil
}

func (b *Bottom) Collect() error {
	for s := range b.collect {
		b.size += s
		b.refer <- s
	}
	close(b.refer)
	return nil
}

func (b *Bottom) Collector() chan int64 {
	return b.collect
}

func (b *Bottom) Level() int {
	return b.level
}

func (b *Bottom) Name() string {
	return filepath.Base(b.path)
}

func (b *Bottom) String() string {
	return fmt.Sprintf("(%s) %s", fmtSize(b.Size(), b.GetUnit()), b.Name())
}

//////////////////////////////
////////// Single ////////////
//////////////////////////////

type Single struct {
	path    string
	size    int64
	walker  Walker
	collect chan int64
}

func (s *Single) Size() int64 {
	return s.size
}

func (s *Single) Spawn(maxLevel int) error {
	s.walker = NewRealWalker(s.path, s.collect)
	go s.walker.Walk()
	return nil
}

func (s *Single) Collect() error {
	for v := range s.collect {
		s.size += v
	}
	return nil
}

func (s *Single) Collector() chan int64 {
	return s.collect
}

func (s *Single) Level() int {
	return 0
}

func NewSingle(path string) *Single {
	s := &Single{path: path}
	s.collect = make(chan int64)
	return s
}

//////////////////////////////
////////// Walker ////////////
//////////////////////////////

type Walker interface {
	Walk()
}