1package tree
2
3import (
4	"bytes"
5)
6
7// Node is a tree node
8type Node struct {
9	Name     string
10	Type     string
11	Template bool
12	Mount    bool
13	Path     string
14	Subtree  *Tree
15}
16
17const (
18	// INF allows to have a full recursion until the leaves of a tree
19	INF = -1
20)
21
22// Nodes is a slice of nodes which can be sorted
23type Nodes []*Node
24
25func (n Nodes) Len() int {
26	return len(n)
27}
28
29func (n Nodes) Less(i, j int) bool {
30	return n[i].Name < n[j].Name
31}
32
33func (n Nodes) Swap(i, j int) {
34	n[i], n[j] = n[j], n[i]
35}
36
37// Equals compares to another node
38func (n Node) Equals(other Node) bool {
39	if n.Name != other.Name {
40		return false
41	}
42	if n.Type != other.Type {
43		return false
44	}
45	if n.Subtree != nil {
46		if other.Subtree == nil {
47			return false
48		}
49		if !n.Subtree.Equals(other.Subtree) {
50			return false
51		}
52	} else {
53		if other.Subtree != nil {
54			return false
55		}
56	}
57	return true
58}
59
60// format returns a pretty printed string of all nodes in and below
61// this node, e.g. ├── baz
62func (n *Node) format(prefix string, last bool, maxDepth, curDepth int) string {
63	if maxDepth > INF && (curDepth > maxDepth+1) {
64		return ""
65	}
66
67	out := bytes.NewBufferString(prefix)
68	// adding either an L or a T, depending if this is the last node
69	// or not
70	if last {
71		_, _ = out.WriteString(symLeaf)
72	} else {
73		_, _ = out.WriteString(symBranch)
74	}
75	// the next levels prefix needs to be extended depending if
76	// this is the last node in a group or not
77	if last {
78		prefix += symEmpty
79	} else {
80		prefix += symVert
81	}
82
83	// any mount will be colored and include the on-disk path
84	switch {
85	case n.Mount:
86		_, _ = out.WriteString(colMount(n.Name + " (" + n.Path + ")"))
87	case n.Type == "dir":
88		_, _ = out.WriteString(colDir(n.Name + sep))
89	default:
90		_, _ = out.WriteString(n.Name)
91	}
92	// mark templates
93	if n.Template {
94		_, _ = out.WriteString(" " + colTpl("(template)"))
95	}
96	// finish this output
97	_, _ = out.WriteString("\n")
98
99	if n.Subtree == nil {
100		return out.String()
101	}
102
103	// let our children format themselves
104	for i, node := range n.Subtree.Nodes {
105		last := i == len(n.Subtree.Nodes)-1
106		_, _ = out.WriteString(node.format(prefix, last, maxDepth, curDepth+1))
107	}
108	return out.String()
109}
110
111// Len returns the length of this subtree
112func (n *Node) Len() int {
113	if n.Type == "file" {
114		return 1
115	}
116	var l int
117	for _, t := range n.Subtree.Nodes {
118		l += t.Len()
119	}
120	return l
121}
122
123func (n *Node) list(prefix string, maxDepth, curDepth int, files bool) []string {
124	if maxDepth >= 0 && curDepth > maxDepth {
125		return nil
126	}
127
128	if prefix != "" {
129		prefix += sep
130	}
131	prefix += n.Name
132
133	// if it's a file and we are looking for files
134	if n.Type == "file" && files {
135		// we return the file
136		return []string{prefix}
137	} else if curDepth == maxDepth && n.Type != "file" {
138		// otherwise if we are "at the bottom" and it's not a file
139		// we return the directory name with a separator at the end
140		return []string{prefix + sep}
141	}
142
143	out := make([]string, 0, n.Len())
144	// if we don't have subitems, then it's a leaf and we return
145	// (notice that this is what ends the recursion when maxDepth is set to -1)
146	if n.Subtree == nil {
147		return out
148	}
149
150	// this is the part that will list the subdirectories on their own line when using the -d option
151	if !files {
152		out = append(out, prefix+sep)
153	}
154
155	// we keep listing the subtree nodes if we haven't exited yet.
156	for _, t := range n.Subtree.Nodes {
157		out = append(out, t.list(prefix, maxDepth, curDepth+1, files)...)
158	}
159	return out
160}
161