Prettify trees

3 files changed, 33 insertions, 11 deletions

diff --git a/cmd/vartan/parse.go b/cmd/vartan/parse.go
index 568a99c..8089c5e 100644
--- a/cmd/vartan/parse.go
+++ b/cmd/vartan/parse.go
@@ -55,7 +55,7 @@ func runParse(cmd *cobra.Command, args []string) (retErr error) {
 		return err
 	}
 	fmt.Printf("Accepted\n")
-	driver.PrintTree(p.AST(), 0)
+	driver.PrintTree(os.Stdout, p.AST())
 
 	return nil
 }
diff --git a/driver/parser.go b/driver/parser.go
index e647a62..7b6df7b 100644
--- a/driver/parser.go
+++ b/driver/parser.go
@@ -14,17 +14,38 @@ type Node struct {
 	Children []*Node
 }
 
-func PrintTree(node *Node, depth int) {
-	for i := 0; i < depth; i++ {
-		fmt.Printf("    ")
+func PrintTree(w io.Writer, node *Node) {
+	printTree(w, node, "", "")
+}
+
+func printTree(w io.Writer, node *Node, ruledLine string, childRuledLinePrefix string) {
+	if node == nil {
+		return
 	}
-	fmt.Printf("%v", node.KindName)
+
 	if node.Text != "" {
-		fmt.Printf(` "%v"`, node.Text)
+		fmt.Fprintf(w, "%v%v %#v\n", ruledLine, node.KindName, node.Text)
+	} else {
+		fmt.Fprintf(w, "%v%v\n", ruledLine, node.KindName)
 	}
-	fmt.Printf("\n")
-	for _, c := range node.Children {
-		PrintTree(c, depth+1)
+
+	num := len(node.Children)
+	for i, child := range node.Children {
+		var line string
+		if num > 1 && i < num-1 {
+			line = "├─ "
+		} else {
+			line = "└─ "
+		}
+
+		var prefix string
+		if i >= num-1 {
+			prefix = "    "
+		} else {
+			prefix = "│  "
+		}
+
+		printTree(w, child, childRuledLinePrefix+line, childRuledLinePrefix+prefix)
 	}
 }
 
diff --git a/driver/parser_test.go b/driver/parser_test.go
index f11a46a..2954ca2 100644
--- a/driver/parser_test.go
+++ b/driver/parser_test.go
@@ -2,6 +2,7 @@ package driver
 
 import (
 	"fmt"
+	"os"
 	"strings"
 	"testing"
 
@@ -271,9 +272,9 @@ foo: "foo";
 			}
 
 			fmt.Println("CST:")
-			PrintTree(p.CST(), 0)
+			PrintTree(os.Stdout, p.CST())
 			fmt.Println("AST:")
-			PrintTree(p.AST(), 0)
+			PrintTree(os.Stdout, p.AST())
 		})
 	}
 }