Merge pull request #11 from shivamMg/bugfix/empty-data

Handle cases of empty data and duplicate nodes

tree/tree.go

@@ -1,6 +1,7 @@
 package tree
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 	"unicode/utf8"
@@ -14,13 +15,24 @@ const (
 	BoxDownRight = "┌"
 	BoxDownHor   = "┬"
 	BoxUpRight   = "└"
-	// Gutter is number of spaces between two adjacent child nodes
+	// Gutter is number of spaces between two adjacent child nodes.
 	Gutter = 2
 )
 
-// Node represents a node in a tree.
+// ErrDuplicateNode indicates that a duplicate Node (node with same hash) was
+// encountered while going through the tree. As of now Sprint/Print and
+// SprintWithError/PrintWithError cannot operate on such trees.
+//
+// This error is returned by SprintWithError/PrintWithError. It's also used
+// in Sprint/Print as error for panic for the same case.
+//
+// FIXME: create internal representation of trees that copies data
+var ErrDuplicateNode = errors.New("duplicate node")
+
+// Node represents a node in a tree. Type that satisfies Node must be a hashable type.
 type Node interface {
-	// Data must return a value representing the node.
+	// Data must return a value representing the node. It is stringified using "%v".
+	// If empty, a space is used.
 	Data() interface{}
 	// Children must return a list of all child nodes of the node.
 	Children() []Node
@@ -58,15 +70,40 @@ func (q *queue) peek() Node {
 	return q.arr[0]
 }
 
-// Print prints the formatted tree to standard output.
+// Print prints the formatted tree to standard output. To handle ErrDuplicateNode use PrintWithError.
 func Print(root Node) {
 	fmt.Print(Sprint(root))
 }
 
-// Sprint returns the formatted tree.
+// Sprint returns the formatted tree. To handle ErrDuplicateNode use SprintWithError.
 func Sprint(root Node) string {
 	parents := map[Node]Node{}
-	setParents(parents, root)
+	if err := setParents(parents, root); err != nil {
+		panic(err)
+	}
+	return sprint(parents, root)
+}
+
+// PrintWithError prints the formatted tree to standard output.
+func PrintWithError(root Node) error {
+	s, err := SprintWithError(root)
+	if err != nil {
+		return err
+	}
+	fmt.Print(s)
+	return nil
+}
+
+// SprintWithError returns the formatted tree.
+func SprintWithError(root Node) (string, error) {
+	parents := map[Node]Node{}
+	if err := setParents(parents, root); err != nil {
+		return "", err
+	}
+	return sprint(parents, root), nil
+}
+
+func sprint(parents map[Node]Node, root Node) string {
 	isLeftMostChild := func(n Node) bool {
 		p, ok := parents[n]
 		if !ok {
@@ -95,7 +132,7 @@ func Sprint(root Node) string {
 			}
 
 			spaces := paddings[n] - covered
-			data := fmt.Sprintf("%v", n.Data())
+			data := safeData(n)
 			nodes += strings.Repeat(" ", spaces) + data
 
 			w := utf8.RuneCountInString(data)
@@ -135,6 +172,17 @@ func Sprint(root Node) string {
 	return s
 }
 
+// safeData always returns non-empty representation of n's data. Empty data
+// messes up tree structure, and ignoring such node will return incomplete
+// tree output (tree without an entire subtree). So it returns a space.
+func safeData(n Node) string {
+	data := fmt.Sprintf("%v", n.Data())
+	if data == "" {
+		return " "
+	}
+	return data
+}
+
 // setPaddings sets left padding (distance of a node from the root)
 // for each node in the tree.
 func setPaddings(paddings map[Node]int, widths map[Node]int, pad int, root Node) {
@@ -147,11 +195,17 @@ func setPaddings(paddings map[Node]int, widths map[Node]int, pad int, root Node)
 
 // setParents sets child-parent relationships for the tree rooted
 // at root.
-func setParents(parents map[Node]Node, root Node) {
+func setParents(parents map[Node]Node, root Node) error {
 	for _, c := range root.Children() {
+		if _, ok := parents[c]; ok {
+			return ErrDuplicateNode
+		}
 		parents[c] = root
-		setParents(parents, c)
+		if err := setParents(parents, c); err != nil {
+			return err
+		}
 	}
+	return nil
 }
 
 // width returns either the sum of widths of it's children or its own
@@ -162,7 +216,7 @@ func width(widths map[Node]int, n Node) int {
 		return w
 	}
 
-	w := utf8.RuneCountInString(fmt.Sprintf("%v", n.Data())) + Gutter
+	w := utf8.RuneCountInString(safeData(n)) + Gutter
 	widths[n] = w
 	if len(n.Children()) == 0 {
 		return w

tree/tree_test.go

@@ -143,3 +143,39 @@ Piano  Orchestra     Opera         Chorus  Rock         Country        Pop
 		t.Errorf("Expected:\n%s\n\nGot:\n%s\n", want, got)
 	}
 }
+
+func TestSprintEmptyData(t *testing.T) {
+	n2, n4, n5 := Node{data: "node2"}, Node{data: "node4"}, Node{data: "node5"}
+	n3 := Node{data: "", c: []*Node{&n4}}
+	n1 := Node{"node1", []*Node{&n2, &n3, &n5}}
+
+	const want = `node1
+├──────┬──────┐
+node2         node5
+       │
+       node4
+`
+
+	got := tree.Sprint(&n1)
+	if got != want {
+		t.Errorf("Expected:\n%s\n\nGot:\n%s\n", want, got)
+	}
+}
+
+func TestSprintAndSprintWithErrorDuplicateNode(t *testing.T) {
+	n2 := Node{data: "a"}
+	n1 := Node{data: "b", c: []*Node{&n2, &n2}}
+
+	_, err := tree.SprintWithError(&n1)
+	if err != tree.ErrDuplicateNode {
+		t.Errorf("Expected:\n%s\n\nGot:\n%s\n", tree.ErrDuplicateNode, err)
+	}
+
+	defer func() {
+		r := recover()
+		if r != tree.ErrDuplicateNode {
+			t.Errorf("Expected:\n%s\n\nGot:\n%s\n", tree.ErrDuplicateNode, r)
+		}
+	}()
+	tree.Sprint(&n1)
+}