Line data Source code
1 : // Copyright 2024 The LevelDB-Go and Pebble Authors. All rights reserved. Use
2 : // of this source code is governed by a BSD-style license that can be found in
3 : // the LICENSE file.
4 :
5 : package treeprinter
6 :
7 : import (
8 : "bytes"
9 : "fmt"
10 : "strings"
11 : )
12 :
13 : var (
14 : edgeLinkChr = rune('│')
15 : edgeMidChr = rune('├')
16 : edgeLastChr = rune('└')
17 : horLineChr = rune('─')
18 : bulletChr = rune('•')
19 : )
20 :
21 : // Node is a handle associated with a specific depth in a tree. See below for
22 : // sample usage.
23 : type Node struct {
24 : tree *tree
25 : level int
26 : }
27 :
28 : // New creates a tree printer and returns a sentinel node reference which
29 : // should be used to add the root. Sample usage:
30 : //
31 : // tp := New()
32 : // root := tp.Child("root")
33 : // root.Child("child-1")
34 : // root.Child("child-2").Child("grandchild\ngrandchild-more-info")
35 : // root.Child("child-3")
36 : //
37 : // fmt.Print(tp.String())
38 : //
39 : // Output:
40 : //
41 : // root
42 : // ├── child-1
43 : // ├── child-2
44 : // │ └── grandchild
45 : // │ grandchild-more-info
46 : // └── child-3
47 : //
48 : // Note that the Child calls can't be rearranged arbitrarily; they have
49 : // to be in the order they need to be displayed (depth-first pre-order).
50 1 : func New() Node {
51 1 : return NewWithStyle(DefaultStyle)
52 1 : }
53 :
54 : // NewWithStyle creates a tree printer like New, permitting customization of
55 : // the style of the resulting tree.
56 1 : func NewWithStyle(style Style) Node {
57 1 : t := &tree{style: style}
58 1 :
59 1 : switch style {
60 1 : case CompactStyle:
61 1 : t.edgeLink = []rune{edgeLinkChr}
62 1 : t.edgeMid = []rune{edgeMidChr, ' '}
63 1 : t.edgeLast = []rune{edgeLastChr, ' '}
64 :
65 1 : case BulletStyle:
66 1 : t.edgeLink = []rune{edgeLinkChr}
67 1 : t.edgeMid = []rune{edgeMidChr, horLineChr, horLineChr, ' '}
68 1 : t.edgeLast = []rune{edgeLastChr, horLineChr, horLineChr, ' '}
69 :
70 1 : default:
71 1 : t.edgeLink = []rune{' ', edgeLinkChr}
72 1 : t.edgeMid = []rune{' ', edgeMidChr, horLineChr, horLineChr, ' '}
73 1 : t.edgeLast = []rune{' ', edgeLastChr, horLineChr, horLineChr, ' '}
74 : }
75 :
76 1 : return Node{
77 1 : tree: t,
78 1 : level: 0,
79 1 : }
80 : }
81 :
82 : // Style is one of the predefined treeprinter styles.
83 : type Style int
84 :
85 : const (
86 : // DefaultStyle is the default style. Example:
87 : //
88 : // foo
89 : // ├── bar1
90 : // │ bar2
91 : // │ └── baz
92 : // └── qux
93 : //
94 : DefaultStyle Style = iota
95 :
96 : // CompactStyle is a compact style, for deeper trees. Example:
97 : //
98 : // foo
99 : // ├ bar1
100 : // │ bar2
101 : // │ └ baz
102 : // └ qux
103 : //
104 : CompactStyle
105 :
106 : // BulletStyle is a style that shows a bullet for each node, and groups any
107 : // other lines under that bullet. Example:
108 : //
109 : // • foo
110 : // │
111 : // ├── • bar1
112 : // │ │ bar2
113 : // │ │
114 : // │ └── • baz
115 : // │
116 : // └── • qux
117 : //
118 : BulletStyle
119 : )
120 :
121 : // tree implements the tree printing machinery.
122 : //
123 : // All Nodes hold a reference to the tree and Node calls result in modification
124 : // of the tree. At any point in time, tree.rows contains the formatted tree that
125 : // was described by the Node calls performed so far.
126 : //
127 : // When new nodes are added, some of the characters of the previous formatted
128 : // tree need to be updated. Here is an example stepping through the state:
129 : //
130 : // API call Rows
131 : //
132 : //
133 : // tp := New() <empty>
134 : //
135 : //
136 : // root := tp.Child("root") root
137 : //
138 : //
139 : // root.Child("child-1") root
140 : // └── child-1
141 : //
142 : //
143 : // c2 := root.Child("child-2") root
144 : // ├── child-1
145 : // └── child-2
146 : //
147 : // Note: here we had to go back up and change └─ into ├─ for child-1.
148 : //
149 : //
150 : // c2.Child("grandchild") root
151 : // ├── child-1
152 : // └── child-2
153 : // └── grandchild
154 : //
155 : //
156 : // root.Child("child-3" root
157 : // ├── child-1
158 : // ├── child-2
159 : // │ └── grandchild
160 : // └── child-3
161 : //
162 : // Note: here we had to go back up and change └─ into ├─ for child-2, and
163 : // add a │ on the grandchild row. In general, we may need to add an
164 : // arbitrary number of vertical bars.
165 : //
166 : // In order to perform these character changes, we maintain information about
167 : // the nodes on the bottom-most path.
168 : type tree struct {
169 : style Style
170 :
171 : // rows maintains the rows accumulated so far, as rune arrays.
172 : rows [][]rune
173 :
174 : // stack contains information pertaining to the nodes on the bottom-most path
175 : // of the tree.
176 : stack []nodeInfo
177 :
178 : edgeLink []rune
179 : edgeMid []rune
180 : edgeLast []rune
181 : }
182 :
183 : type nodeInfo struct {
184 : // firstChildConnectRow is the index (in tree.rows) of the row up to which we
185 : // have to connect the first child of this node.
186 : firstChildConnectRow int
187 :
188 : // nextSiblingConnectRow is the index (in tree.rows) of the row up to which we
189 : // have to connect the next sibling of this node. Typically this is the same
190 : // with firstChildConnectRow, except when the node has multiple rows. For
191 : // example:
192 : //
193 : // foo
194 : // └── bar1 <---- nextSiblingConnectRow
195 : // bar2 <---- firstChildConnectRow
196 : //
197 : // firstChildConnectRow is used when adding "baz", nextSiblingConnectRow
198 : // is used when adding "qux":
199 : // foo
200 : // ├── bar1
201 : // │ bar2
202 : // │ └── baz
203 : // └── qux
204 : //
205 : nextSiblingConnectRow int
206 : }
207 :
208 : // set copies the string of runes into a given row, at a specific position. The
209 : // row is extended with spaces if needed.
210 1 : func (t *tree) set(rowIdx int, colIdx int, what []rune) {
211 1 : // Extend the line if necessary.
212 1 : for len(t.rows[rowIdx]) < colIdx+len(what) {
213 1 : t.rows[rowIdx] = append(t.rows[rowIdx], ' ')
214 1 : }
215 1 : copy(t.rows[rowIdx][colIdx:], what)
216 : }
217 :
218 : // addRow adds a row with a given text, with the proper indentation for the
219 : // given level.
220 1 : func (t *tree) addRow(level int, text string) (rowIdx int) {
221 1 : runes := []rune(text)
222 1 : // Each level indents by this much.
223 1 : k := len(t.edgeLast)
224 1 : indent := level * k
225 1 : row := make([]rune, indent+len(runes))
226 1 : for i := 0; i < indent; i++ {
227 1 : row[i] = ' '
228 1 : }
229 1 : copy(row[indent:], runes)
230 1 : t.rows = append(t.rows, row)
231 1 : return len(t.rows) - 1
232 : }
233 :
234 : // Childf adds a node as a child of the given node.
235 1 : func (n Node) Childf(format string, args ...interface{}) Node {
236 1 : return n.Child(fmt.Sprintf(format, args...))
237 1 : }
238 :
239 : // Child adds a node as a child of the given node. Multi-line strings are
240 : // supported with appropriate indentation.
241 1 : func (n Node) Child(text string) Node {
242 1 : if strings.ContainsRune(text, '\n') {
243 1 : splitLines := strings.Split(text, "\n")
244 1 : node := n.childLine(splitLines[0])
245 1 : for _, l := range splitLines[1:] {
246 1 : node.AddLine(l)
247 1 : }
248 1 : return node
249 : }
250 1 : return n.childLine(text)
251 : }
252 :
253 : // AddLine adds a new line to a node without an edge.
254 1 : func (n Node) AddLine(text string) {
255 1 : t := n.tree
256 1 : if t.style == BulletStyle {
257 1 : text = " " + text
258 1 : }
259 1 : rowIdx := t.addRow(n.level-1, text)
260 1 : if t.style != BulletStyle {
261 1 : t.stack[n.level-1].firstChildConnectRow = rowIdx
262 1 : }
263 : }
264 :
265 : // childLine adds a node as a child of the given node.
266 1 : func (n Node) childLine(text string) Node {
267 1 : t := n.tree
268 1 : if t.style == BulletStyle {
269 1 : text = fmt.Sprintf("%c %s", bulletChr, text)
270 1 : if n.level > 0 {
271 1 : n.AddEmptyLine()
272 1 : }
273 : }
274 1 : rowIdx := t.addRow(n.level, text)
275 1 : edgePos := (n.level - 1) * len(t.edgeLast)
276 1 : if n.level == 0 {
277 1 : // Case 1: root.
278 1 : if len(t.stack) != 0 {
279 0 : panic("multiple root nodes")
280 : }
281 1 : } else if len(t.stack) <= n.level {
282 1 : // Case 2: first child. Connect to parent.
283 1 : if len(t.stack) != n.level {
284 0 : panic("misuse of node")
285 : }
286 1 : parentRow := t.stack[n.level-1].firstChildConnectRow
287 1 : for i := parentRow + 1; i < rowIdx; i++ {
288 1 : t.set(i, edgePos, t.edgeLink)
289 1 : }
290 1 : t.set(rowIdx, edgePos, t.edgeLast)
291 1 : } else {
292 1 : // Case 3: non-first child. Connect to sibling.
293 1 : siblingRow := t.stack[n.level].nextSiblingConnectRow
294 1 : t.set(siblingRow, edgePos, t.edgeMid)
295 1 : for i := siblingRow + 1; i < rowIdx; i++ {
296 1 : t.set(i, edgePos, t.edgeLink)
297 1 : }
298 1 : t.set(rowIdx, edgePos, t.edgeLast)
299 1 : // Update the nextSiblingConnectRow.
300 1 : t.stack = t.stack[:n.level]
301 : }
302 :
303 1 : t.stack = append(t.stack, nodeInfo{
304 1 : firstChildConnectRow: rowIdx,
305 1 : nextSiblingConnectRow: rowIdx,
306 1 : })
307 1 :
308 1 : // Return a TreePrinter that can be used for children of this node.
309 1 : return Node{
310 1 : tree: t,
311 1 : level: n.level + 1,
312 1 : }
313 : }
314 :
315 : // AddEmptyLine adds an empty line to the output; used to introduce vertical
316 : // spacing as needed.
317 1 : func (n Node) AddEmptyLine() {
318 1 : n.tree.rows = append(n.tree.rows, []rune{})
319 1 : }
320 :
321 : // FormattedRows returns the formatted rows. Can only be called on the result of
322 : // treeprinter.New.
323 1 : func (n Node) FormattedRows() []string {
324 1 : if n.level != 0 {
325 0 : panic("Only the root can be stringified")
326 : }
327 1 : res := make([]string, len(n.tree.rows))
328 1 : for i, r := range n.tree.rows {
329 1 : res[i] = string(r)
330 1 : }
331 1 : return res
332 : }
333 :
334 1 : func (n Node) String() string {
335 1 : if n.level != 0 {
336 0 : panic("Only the root can be stringified")
337 : }
338 1 : var buf bytes.Buffer
339 1 : for _, r := range n.tree.rows {
340 1 : buf.WriteString(string(r))
341 1 : buf.WriteByte('\n')
342 1 : }
343 1 : return buf.String()
344 : }
|
