Line data    Source code

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