From 29c21cb78d9161cbe8f14ee0879f125363eb923c Mon Sep 17 00:00:00 2001
From: Grzegorz Dlugoszewski <g.dlugoszewski@gmail.com>
Date: Sat, 6 Jun 2020 16:54:03 +0200
Subject: [PATCH] Improve the implementation of the smart tree printer

It now correctly indents lines with branches status
---
 pkg/list.go      |  2 +-
 pkg/tree.go      | 45 +++++++++++++++++++++++++++++++++++++++------
 pkg/tree_test.go |  8 +++++---
 3 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/pkg/list.go b/pkg/list.go
index 9410842..535142c 100644
--- a/pkg/list.go
+++ b/pkg/list.go
@@ -106,7 +106,7 @@ func PrintRepos(repos []*Repo) {
 	root := viper.GetString(KeyReposRoot)
 
 	tree := BuildTree(root, repos)
-	fmt.Println(RenderTree(tree))
+	fmt.Println(RenderSmartTree(tree))
 }
 
 const (
diff --git a/pkg/tree.go b/pkg/tree.go
index 9808884..e0df7e7 100644
--- a/pkg/tree.go
+++ b/pkg/tree.go
@@ -86,7 +86,9 @@ func BuildTree(root string, repos []*Repo) *Node {
 	return tree
 }
 
-// RenderTree returns a string representation of repos tree.
+// RenderSmartTree returns a string representation of repos tree.
+// It's "smart" because it automatically folds branches which only have a single child and indents branches with many children.
+//
 // It recursively traverses the tree and prints its nodes.
 // If a node contains multiple children, they are be printed in new lines and indented.
 // If a node contains only a single child, it is printed in the same line using path separator.
@@ -106,12 +108,35 @@ func BuildTree(root string, repos []*Repo) *Node {
 //           repo2
 //       another/repo
 //
-func RenderTree(node *Node) string {
+func RenderSmartTree(node *Node) string {
 	if node.children == nil {
 		// If node is a leaf, print repo name and its status and finish processing this node.
-		return node.val + " " + renderWorktreeStatus(node.repo)
+		value := node.val
+
+		// TODO: Ugly
+		// If this is called from tests the repo will be nil and we should return just the name without the status.
+		if node.repo.repo == nil {
+			return value
+		}
+
+		value += " " + renderWorktreeStatus(node.repo)
+
+		// Print the status of each branch on a new line, indented to match the position of the current branch name.
+		indent := "\n" + strings.Repeat(" ", length+len(node.val))
+		for _, branch := range node.repo.Status.Branches {
+			// Don't print the status of the current branch. It was already printed above.
+			if branch.Name == node.repo.Status.CurrentBranch {
+				continue
+			}
+
+			value += indent + renderBranchStatus(branch)
+		}
+
+		return value
 	}
 
+	val := node.val + string(filepath.Separator)
+
 	shift := ""
 	if node.parent == nil {
 		// If node is a root, print its children on a new line without indentation.
@@ -121,17 +146,25 @@ func RenderTree(node *Node) string {
 		// Setting node's depth to the same as parent's ensures that its children will be indented only once even if
 		// node's path has multiple levels above.
 		node.depth = node.parent.depth
+
+		length += len(val)
 	} else {
 		// If node has multiple children, print each of them on a new line
 		// and indent them once relative to the parent
 		node.depth = node.parent.depth + 1
-		shift = "\n" + strings.Repeat("\t", node.depth)
+		shift = "\n" + strings.Repeat("    ", node.depth)
+		length = 0
 	}
 
-	val := node.val + string(filepath.Separator)
 	for _, child := range node.children {
-		val += shift + RenderTree(child)
+		length += len(shift)
+		val += shift + RenderSmartTree(child)
+		length = 0
 	}
 
 	return val
 }
+
+// lenght is the size (number of chars) of the currently processed line.
+// It's used to correctly indent the lines with branches status.
+var length int
diff --git a/pkg/tree_test.go b/pkg/tree_test.go
index 1da16df..a31d384 100644
--- a/pkg/tree_test.go
+++ b/pkg/tree_test.go
@@ -2,6 +2,7 @@ package pkg
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 )
 
@@ -91,14 +92,15 @@ gitlab.com/
 	for i, test := range tests {
 		var repos []*Repo
 		for _, path := range test.paths {
-			repos = append(repos, &Repo{path: path})
+			repos = append(repos, newRepo(nil, path)) //&Repo{path: path})
 		}
 
 		tree := BuildTree("root", repos)
 		// Leading and trailing newlines are added to test cases for readability. We also need to add them to the rendering result.
-		got := fmt.Sprintf("\n%s\n", RenderTree(tree))
+		got := fmt.Sprintf("\n%s\n", RenderSmartTree(tree))
 
-		if got != test.want {
+		// Rendered tree uses spaces for indentation but the test cases use tabs.
+		if got != strings.ReplaceAll(test.want, "\t", "    ") {
 			t.Errorf("Failed test case %d, got: %+v; want: %+v", i, got, test.want)
 		}
 	}