ruspa/tree/dynamic_tree.go

package tree

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

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

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
	AddCollector(chan int64)
	Level() int
	Name() string
	Complete() bool
}

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

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

func NewTop(path string) *Top {
	t := &Top{path: path, level: 0, unit: KB}
	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
			var collect chan int64
			path := filepath.Join(t.path, info.Name())
			if isLastLevel(t, maxLevel) || !AnyDirectoryDownThere(path) {
				child, collect = t.newBottom(path)
			} else {
				child, collect = t.newIntermediate(path)
			}
			t.AddCollector(collect)
			go child.Spawn(maxLevel)
		case !mode.IsDir() && mode.IsRegular():
			t.size += info.Size()
		}
	}

	go t.Collect()
	return nil
}

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

func (t *Top) AddCollector(collect chan int64) {
	t.collect = append(t.collect, collect)
}

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

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

func (t *Top) Complete() bool {
	for _, child := range t.tree {
		if !child.Complete() {
			return false
		}
	}
	return true
}

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, chan int64) {
	switch {
	case kind == IntermediateType:
		i, collect := NewIntermediate(path, t)
		t.addToTree(i)
		return i, collect
	case kind == BottomType:
		b, collect := NewBottom(path, t)
		t.addToTree(b)
		return b, collect
	default:
		panic(ErrUnknownType)
	}
}

func (t *Top) newIntermediate(path string) (*Intermediate, chan int64) {
	child, collect := t.newChild(path, IntermediateType)
	return child.(*Intermediate), collect
}

func (t *Top) newBottom(path string) (*Bottom, chan int64) {
	child, collect := t.newChild(path, BottomType)
	return child.(*Bottom), collect
}

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

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

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

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
			var collect chan int64
			path := filepath.Join(i.path, info.Name())
			if isLastLevel(i, maxLevel) || !AnyDirectoryDownThere(path) {
				child, collect = i.newBottom(path)
			} else {
				child, collect = i.newIntermediate(path)
			}
			i.AddCollector(collect)
			go child.Spawn(maxLevel)
		case !mode.IsDir() && mode.IsRegular():
			i.size += info.Size()
			i.refer <- info.Size()
		}
	}

	go i.Collect()
	return nil
}

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

func (i *Intermediate) AddCollector(collect chan int64) {
	i.collect = append(i.collect, collect)
}

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

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

func (i *Intermediate) Complete() bool {
	for _, child := range i.tree {
		if !child.Complete() {
			return false
		}
	}
	return true
}

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, chan int64) {
	switch {
	case kind == IntermediateType:
		c, collect := NewIntermediate(path, i)
		i.addToTree(c)
		return c, collect
	case kind == BottomType:
		b, collect := NewBottom(path, i)
		i.addToTree(b)
		return b, collect
	default:
		panic(ErrUnknownType)
	}
}

func (i *Intermediate) newIntermediate(path string) (*Intermediate, chan int64) {
	child, collect := i.newChild(path, IntermediateType)
	return child.(*Intermediate), collect
}

func (i *Intermediate) newBottom(path string) (*Bottom, chan int64) {
	child, collect := i.newChild(path, BottomType)
	return child.(*Bottom), collect
}

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

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

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

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 {
	collect := make(chan int64)
	b.AddCollector(collect)
	b.walker = NewRealWalker(b.path, collect)
	go b.walker.Walk()
	go b.Collect()
	return nil
}

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

func (b *Bottom) AddCollector(collect chan int64) {
	b.collect = collect
}

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

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

func (b *Bottom) Complete() bool {
	return b.complete
}

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
	unit     ByteSize
	walker   Walker
	collect  chan int64
	complete bool
}

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

func (s *Single) SetUnit(unit string) {
	s.unit = setUnit(unit)
}

func (s *Single) GetUnit() ByteSize {
	return s.unit
}

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()
	go s.Collect()
	return nil
}

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

func (s *Single) AddCollector(collect chan int64) {
	s.collect = collect
}

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

func (s *Single) Name() string {
	return filepath.Base(s.path)
}

func (s *Single) Complete() bool {
	return s.complete
}

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

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

type Walker interface {
	Walk()
}