refactor

3 files changed, 28 insertions, 27 deletions

diff --git a/pkg/commands/models/status_line.go b/pkg/commands/models/file_change_interface.go
index 5413b076f..5413b076f 100644
--- a/pkg/commands/models/status_line.go
+++ b/pkg/commands/models/file_change_interface.go
diff --git a/pkg/commands/models/status_line_node.go b/pkg/commands/models/file_change_node.go
index 9e0f241f3..d265b3f05 100644
--- a/pkg/commands/models/status_line_node.go
+++ b/pkg/commands/models/file_change_node.go
@@ -9,10 +9,11 @@ import (
 )
 
 type FileChangeNode struct {
-	Children  []*FileChangeNode
-	File      *File
-	Path      string // e.g. '/path/to/mydir'
-	Collapsed bool
+	Children         []*FileChangeNode
+	File             *File
+	Path             string // e.g. '/path/to/mydir'
+	Collapsed        bool
+	CompressionLevel int // equal to the number of forward slashes you'll see in the path when it's rendered
 }
 
 func (s *FileChangeNode) GetHasUnstagedChanges() bool {
@@ -186,8 +187,10 @@ func (s *FileChangeNode) compressAux() *FileChangeNode {
 
 	for i := range s.Children {
 		for s.Children[i].HasExactlyOneChild() {
+			prevCompressionLevel := s.Children[i].CompressionLevel
 			grandchild := s.Children[i].Children[0]
 			s.Children[i] = grandchild
+			s.Children[i].CompressionLevel = prevCompressionLevel + 1
 		}
 	}
 
@@ -241,6 +244,19 @@ func (s *FileChangeNode) ForEachFile(cb func(*File) error) error {
 	return nil
 }
 
+func (s *FileChangeNode) GetLeaves() []*FileChangeNode {
+	if s.IsLeaf() {
+		return []*FileChangeNode{s}
+	}
+
+	output := []*FileChangeNode{}
+	for _, child := range s.Children {
+		output = append(output, child.GetLeaves()...)
+	}
+
+	return output
+}
+
 func (s *FileChangeNode) NameAtDepth(depth int) string {
 	splitName := strings.Split(s.Path, string(os.PathSeparator))
 	name := filepath.Join(splitName[depth:]...)
diff --git a/pkg/commands/models/status_line_node_test.go b/pkg/commands/models/file_change_node_test.go
index 0f0f95ce5..8190230e3 100644
--- a/pkg/commands/models/status_line_node_test.go
+++ b/pkg/commands/models/file_change_node_test.go
@@ -20,61 +20,53 @@ func TestCompress(t *testing.T) {
 		{
 			name: "leaf node",
 			root: &FileChangeNode{
-				Name: "",
+				Path: "",
 				Children: []*FileChangeNode{
-					{File: &File{Name: "test", ShortStatus: " M", HasStagedChanges: true}, Name: "test"},
+					{File: &File{Name: "test", ShortStatus: " M", HasStagedChanges: true}, Path: "test"},
 				},
 			},
 			expected: &FileChangeNode{
-				Name: "",
+				Path: "",
 				Children: []*FileChangeNode{
-					{File: &File{Name: "test", ShortStatus: " M", HasStagedChanges: true}, Name: "test"},
+					{File: &File{Name: "test", ShortStatus: " M", HasStagedChanges: true}, Path: "test"},
 				},
 			},
 		},
 		{
 			name: "big example",
 			root: &FileChangeNode{
-				Name: "",
+				Path: "",
 				Children: []*FileChangeNode{
 					{
-						Name: "dir1",
 						Path: "dir1",
 						Children: []*FileChangeNode{
 							{
 								File: &File{Name: "file2", ShortStatus: "M ", HasUnstagedChanges: true},
-								Name: "file2",
 								Path: "dir1/file2",
 							},
 						},
 					},
 					{
-						Name: "dir2",
 						Path: "dir2",
 						Children: []*FileChangeNode{
 							{
 								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",
 							},
 						},
 					},
 					{
-						Name: "dir3",
 						Path: "dir3",
 						Children: []*FileChangeNode{
 							{
-								Name: "dir3-1",
 								Path: "dir3/dir3-1",
 								Children: []*FileChangeNode{
 									{
 										File: &File{Name: "file5", ShortStatus: "M ", HasUnstagedChanges: true},
-										Name: "file5",
 										Path: "dir3/dir3-1/file5",
 									},
 								},
@@ -83,43 +75,36 @@ func TestCompress(t *testing.T) {
 					},
 					{
 						File: &File{Name: "file1", ShortStatus: "M ", HasUnstagedChanges: true},
-						Name: "file1",
 						Path: "file1",
 					},
 				},
 			},
 			expected: &FileChangeNode{
-				Name: "",
+				Path: "",
 				Children: []*FileChangeNode{
 					{
-						Name: "dir1/file2",
-						File: &File{Name: "file2", ShortStatus: "M ", HasUnstagedChanges: true},
 						Path: "dir1/file2",
+						File: &File{Name: "file2", ShortStatus: "M ", HasUnstagedChanges: true},
 					},
 					{
-						Name: "dir2",
 						Path: "dir2",
 						Children: []*FileChangeNode{
 							{
 								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",
 							},
 						},
 					},
 					{
-						Name: "dir3/dir3-1/file5",
-						File: &File{Name: "file5", ShortStatus: "M ", HasUnstagedChanges: true},
 						Path: "dir3/dir3-1/file5",
+						File: &File{Name: "file5", ShortStatus: "M ", HasUnstagedChanges: true},
 					},
 					{
 						File: &File{Name: "file1", ShortStatus: "M ", HasUnstagedChanges: true},
-						Name: "file1",
 						Path: "file1",
 					},
 				},