combining nodes when only one child exists

2 files changed, 84 insertions, 12 deletions

diff --git a/pkg/commands/models/status_line_node.go b/pkg/commands/models/status_line_node.go
index 92f0bba03..b33e9e731 100644
--- a/pkg/commands/models/status_line_node.go
+++ b/pkg/commands/models/status_line_node.go
@@ -1,6 +1,7 @@
 package models
 
 import (
+	"fmt"
 	"sort"
 )
 
@@ -106,11 +107,11 @@ func (s *StatusLineNode) Flatten() []*StatusLineNode {
 	return arr
 }
 
-func (s *StatusLineNode) SortTree() {
+func (s *StatusLineNode) Sort() {
 	s.sortChildren()
 
 	for _, child := range s.Children {
-		child.SortTree()
+		child.Sort()
 	}
 }
 
@@ -149,3 +150,35 @@ func (s *StatusLineNode) GetIsTracked() bool {
 func (s *StatusLineNode) GetPath() string {
 	return s.Path
 }
+
+func (s *StatusLineNode) HasExactlyOneChild() bool {
+	return len(s.Children) == 1
+}
+
+func (s *StatusLineNode) Compress() {
+	if s == nil {
+		return
+	}
+
+	s.compressAux()
+}
+
+func (s *StatusLineNode) compressAux() *StatusLineNode {
+	if s.IsLeaf() {
+		return s
+	}
+
+	for i, child := range s.Children {
+		if child.HasExactlyOneChild() {
+			grandchild := child.Children[0]
+			grandchild.Name = fmt.Sprintf("%s/%s", child.Name, grandchild.Name)
+			s.Children[i] = grandchild
+		}
+	}
+
+	for i, child := range s.Children {
+		s.Children[i] = child.compressAux()
+	}
+
+	return s
+}
diff --git a/pkg/commands/models/status_line_node_test.go b/pkg/commands/models/status_line_node_test.go
index f8114fcf1..c68199c18 100644
--- a/pkg/commands/models/status_line_node_test.go
+++ b/pkg/commands/models/status_line_node_test.go
@@ -6,16 +6,16 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestRender(t *testing.T) {
+func TestCompress(t *testing.T) {
 	scenarios := []struct {
 		name     string
 		root     *StatusLineNode
-		expected []string
+		expected *StatusLineNode
 	}{
 		{
 			name:     "nil node",
 			root:     nil,
-			expected: []string{},
+			expected: nil,
 		},
 		{
 			name: "leaf node",
@@ -25,7 +25,12 @@ func TestRender(t *testing.T) {
 					{File: &File{Name: "test", ShortStatus: " M", HasStagedChanges: true}, Name: "test"},
 				},
 			},
-			expected: []string{" M test"},
+			expected: &StatusLineNode{
+				Name: "",
+				Children: []*StatusLineNode{
+					{File: &File{Name: "test", ShortStatus: " M", HasStagedChanges: true}, Name: "test"},
+				},
+			},
 		},
 		{
 			name: "big example",
@@ -33,44 +38,78 @@ func TestRender(t *testing.T) {
 				Name: "",
 				Children: []*StatusLineNode{
 					{
-						Name:      "dir1",
-						Collapsed: true,
+						Name: "dir1",
+						Path: "dir1",
 						Children: []*StatusLineNode{
 							{
 								File: &File{Name: "file2", ShortStatus: "M ", HasUnstagedChanges: true},
 								Name: "file2",
+								Path: "dir1/file2",
 							},
 						},
 					},
 					{
 						Name: "dir2",
+						Path: "dir2",
 						Children: []*StatusLineNode{
 							{
 								File: &File{Name: "file3", ShortStatus: " M", HasStagedChanges: true},
 								Name: "file3",
+								Path: "dir2/file3",
 							},
 							{
 								File: &File{Name: "file4", ShortStatus: "M ", HasUnstagedChanges: true},
 								Name: "file4",
+								Path: "dir2/file4",
 							},
 						},
 					},
 					{
 						File: &File{Name: "file1", ShortStatus: "M ", HasUnstagedChanges: true},
 						Name: "file1",
+						Path: "file1",
+					},
+				},
+			},
+			expected: &StatusLineNode{
+				Name: "",
+				Children: []*StatusLineNode{
+					{
+						Name: "dir1/file2",
+						File: &File{Name: "file2", ShortStatus: "M ", HasUnstagedChanges: true},
+						Path: "dir1/file2",
+					},
+					{
+						Name: "dir2",
+						Path: "dir2",
+						Children: []*StatusLineNode{
+							{
+								File: &File{Name: "file3", ShortStatus: " M", HasStagedChanges: true},
+								Name: "file3",
+								Path: "dir2/file3",
+							},
+							{
+								File: &File{Name: "file4", ShortStatus: "M ", HasUnstagedChanges: true},
+								Name: "file4",
+								Path: "dir2/file4",
+							},
+						},
+					},
+					{
+						File: &File{Name: "file1", ShortStatus: "M ", HasUnstagedChanges: true},
+						Name: "file1",
+						Path: "file1",
 					},
 				},
 			},
-
-			expected: []string{"M  dir1 ►", "MM dir2 ▼", "   M file3", "  M  file4", "M  file1"},
 		},
 	}
 
 	for _, s := range scenarios {
 		s := s
 		t.Run(s.name, func(t *testing.T) {
-			result := s.root.Render()[1:]
-			assert.EqualValues(t, s.expected, result)
+			s.root.Compress()
+			assert.EqualValues(t, s.expected, s.root)
 		})
 	}
 }