refactor to support commit file tree

3 files changed, 439 insertions, 191 deletions

diff --git a/pkg/gui/filetree/commit_file_change_node.go b/pkg/gui/filetree/commit_file_change_node.go
new file mode 100644
index 000000000..c1f99d937
--- /dev/null
+++ b/pkg/gui/filetree/commit_file_change_node.go
@@ -0,0 +1,147 @@
+package filetree
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+)
+
+type CommitFileChangeNode struct {
+	Children         []*CommitFileChangeNode
+	File             *models.CommitFile
+	Path             string // e.g. '/path/to/mydir'
+	CompressionLevel int    // equal to the number of forward slashes you'll see in the path when it's rendered in tree mode
+}
+
+// methods satisfying ListItem interface
+
+func (s *CommitFileChangeNode) ID() string {
+	return s.GetPath()
+}
+
+func (s *CommitFileChangeNode) Description() string {
+	return s.GetPath()
+}
+
+// methods satisfying INode interface
+
+func (s *CommitFileChangeNode) IsLeaf() bool {
+	return s.File != nil
+}
+
+func (s *CommitFileChangeNode) GetPath() string {
+	return s.Path
+}
+
+func (s *CommitFileChangeNode) GetChildren() []INode {
+	result := make([]INode, len(s.Children))
+	for i, child := range s.Children {
+		result[i] = child
+	}
+
+	return result
+}
+
+func (s *CommitFileChangeNode) SetChildren(children []INode) {
+	castChildren := make([]*CommitFileChangeNode, len(children))
+	for i, child := range children {
+		castChildren[i] = child.(*CommitFileChangeNode)
+	}
+
+	s.Children = castChildren
+}
+
+func (s *CommitFileChangeNode) GetCompressionLevel() int {
+	return s.CompressionLevel
+}
+
+func (s *CommitFileChangeNode) SetCompressionLevel(level int) {
+	s.CompressionLevel = level
+}
+
+// methods utilising generic functions for INodes
+
+func (s *CommitFileChangeNode) Sort() {
+	sortNode(s)
+}
+
+func (s *CommitFileChangeNode) ForEachFile(cb func(*models.CommitFile) error) error {
+	return forEachLeaf(s, func(n INode) error {
+		castNode := n.(*CommitFileChangeNode)
+		return cb(castNode.File)
+	})
+}
+
+func (s *CommitFileChangeNode) Any(test func(node *CommitFileChangeNode) bool) bool {
+	return any(s, func(n INode) bool {
+		castNode := n.(*CommitFileChangeNode)
+		return test(castNode)
+	})
+}
+
+func (n *CommitFileChangeNode) Flatten(collapsedPaths map[string]bool) []*CommitFileChangeNode {
+	results := flatten(n, collapsedPaths)
+	nodes := make([]*CommitFileChangeNode, len(results))
+	for i, result := range results {
+		nodes[i] = result.(*CommitFileChangeNode)
+	}
+
+	return nodes
+}
+
+func (node *CommitFileChangeNode) GetNodeAtIndex(index int, collapsedPaths map[string]bool) *CommitFileChangeNode {
+	return getNodeAtIndex(node, index, collapsedPaths).(*CommitFileChangeNode)
+}
+
+func (node *CommitFileChangeNode) GetIndexForPath(path string, collapsedPaths map[string]bool) (int, bool) {
+	return getIndexForPath(node, path, collapsedPaths)
+}
+
+func (node *CommitFileChangeNode) Size(collapsedPaths map[string]bool) int {
+	return size(node, collapsedPaths)
+}
+
+func (s *CommitFileChangeNode) Compress() {
+	// with these functions I try to only have type conversion code on the actual struct,
+	// but comparing interface values to nil is fraught with danger so I'm duplicating
+	// that code here.
+	if s == nil {
+		return
+	}
+
+	compressAux(s)
+}
+
+// This ignores the root
+func (node *CommitFileChangeNode) GetPathsMatching(test func(*CommitFileChangeNode) bool) []string {
+	return getPathsMatching(node, func(n INode) bool {
+		return test(n.(*CommitFileChangeNode))
+	})
+}
+
+func (s *CommitFileChangeNode) GetLeaves() []*CommitFileChangeNode {
+	leaves := getLeaves(s)
+	castLeaves := make([]*CommitFileChangeNode, len(leaves))
+	for i := range leaves {
+		castLeaves[i] = leaves[i].(*CommitFileChangeNode)
+	}
+
+	return castLeaves
+}
+
+// extra methods
+
+func (s *CommitFileChangeNode) AnyFile(test func(file *models.CommitFile) bool) bool {
+	return s.Any(func(node *CommitFileChangeNode) bool {
+		return node.IsLeaf() && test(node.File)
+	})
+}
+
+func (s *CommitFileChangeNode) NameAtDepth(depth int) string {
+	splitName := strings.Split(s.Path, string(os.PathSeparator))
+	name := filepath.Join(splitName[depth:]...)
+
+	return name
+}
diff --git a/pkg/gui/filetree/file_change_node.go b/pkg/gui/filetree/file_change_node.go
index 0515f976c..a99a47ac1 100644
--- a/pkg/gui/filetree/file_change_node.go
+++ b/pkg/gui/filetree/file_change_node.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"sort"
 	"strings"
 
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
@@ -17,257 +16,163 @@ type FileChangeNode struct {
 	CompressionLevel int    // equal to the number of forward slashes you'll see in the path when it's rendered in tree mode
 }
 
-func (s *FileChangeNode) GetHasUnstagedChanges() bool {
-	return s.AnyFile(func(file *models.File) bool { return file.HasUnstagedChanges })
-}
+// methods satisfying ListItem interface
 
-func (s *FileChangeNode) GetHasStagedChanges() bool {
-	return s.AnyFile(func(file *models.File) bool { return file.HasStagedChanges })
+func (s *FileChangeNode) ID() string {
+	return s.GetPath()
 }
 
-func (s *FileChangeNode) GetHasInlineMergeConflicts() bool {
-	return s.AnyFile(func(file *models.File) bool { return file.HasInlineMergeConflicts })
+func (s *FileChangeNode) Description() string {
+	return s.GetPath()
 }
 
-func (s *FileChangeNode) GetIsTracked() bool {
-	return s.AnyFile(func(file *models.File) bool { return file.Tracked })
-}
+// methods satisfying INode interface
 
-func (s *FileChangeNode) AnyFile(test func(file *models.File) bool) bool {
-	return s.Any(func(node *FileChangeNode) bool {
-		return node.IsLeaf() && test(node.File)
-	})
+func (s *FileChangeNode) IsLeaf() bool {
+	return s.File != nil
 }
 
-func (s *FileChangeNode) ForEachFile(cb func(*models.File) error) error {
-	if s.File != nil {
-		if err := cb(s.File); err != nil {
-			return err
-		}
-	}
-
-	for _, child := range s.Children {
-		if err := child.ForEachFile(cb); err != nil {
-			return err
-		}
-	}
-
-	return nil
+func (s *FileChangeNode) GetPath() string {
+	return s.Path
 }
 
-func (s *FileChangeNode) NameAtDepth(depth int) string {
-	splitName := strings.Split(s.Path, string(os.PathSeparator))
-	name := filepath.Join(splitName[depth:]...)
-
-	if s.File != nil && s.File.IsRename() {
-		splitPrevName := strings.Split(s.File.PreviousName, string(os.PathSeparator))
-
-		prevName := s.File.PreviousName
-		// if the file has just been renamed inside the same directory, we can shave off
-		// the prefix for the previous path too. Otherwise we'll keep it unchanged
-		sameParentDir := len(splitName) == len(splitPrevName) && filepath.Join(splitName[0:depth]...) == filepath.Join(splitPrevName[0:depth]...)
-		if sameParentDir {
-			prevName = filepath.Join(splitPrevName[depth:]...)
-		}
-
-		return fmt.Sprintf("%s%s%s", prevName, " → ", name)
+func (s *FileChangeNode) GetChildren() []INode {
+	result := make([]INode, len(s.Children))
+	for i, child := range s.Children {
+		result[i] = child
 	}
 
-	return name
+	return result
 }
 
-func (s *FileChangeNode) Any(test func(node *FileChangeNode) bool) bool {
-	if test(s) {
-		return true
+func (s *FileChangeNode) SetChildren(children []INode) {
+	castChildren := make([]*FileChangeNode, len(children))
+	for i, child := range children {
+		castChildren[i] = child.(*FileChangeNode)
 	}
 
-	for _, child := range s.Children {
-		if child.Any(test) {
-			return true
-		}
-	}
-
-	return false
+	s.Children = castChildren
 }
 
-func (s *FileChangeNode) GetNodeAtIndex(index int, collapsedPaths map[string]bool) *FileChangeNode {
-	node, _ := s.getNodeAtIndexAux(index, collapsedPaths)
-
-	return node
+func (s *FileChangeNode) GetCompressionLevel() int {
+	return s.CompressionLevel
 }
 
-func (s *FileChangeNode) getNodeAtIndexAux(index int, collapsedPaths map[string]bool) (*FileChangeNode, int) {
-	offset := 1
-
-	if index == 0 {
-		return s, offset
-	}
-
-	if !collapsedPaths[s.GetPath()] {
-		for _, child := range s.Children {
-			node, offsetChange := child.getNodeAtIndexAux(index-offset, collapsedPaths)
-			offset += offsetChange
-			if node != nil {
-				return node, offset
-			}
-		}
-	}
-
-	return nil, offset
+func (s *FileChangeNode) SetCompressionLevel(level int) {
+	s.CompressionLevel = level
 }
 
-func (s *FileChangeNode) GetIndexForPath(path string, collapsedPaths map[string]bool) (int, bool) {
-	return s.getIndexForPathAux(path, collapsedPaths)
-}
+// methods utilising generic functions for INodes
 
-func (s *FileChangeNode) getIndexForPathAux(path string, collapsedPaths map[string]bool) (int, bool) {
-	offset := 0
-
-	if s.Path == path {
-		return offset, true
-	}
-
-	if !collapsedPaths[s.GetPath()] {
-		for _, child := range s.Children {
-			offsetChange, found := child.getIndexForPathAux(path, collapsedPaths)
-			offset += offsetChange + 1
-			if found {
-				return offset, true
-			}
-		}
-	}
-
-	return offset, false
+func (s *FileChangeNode) Sort() {
+	sortNode(s)
 }
 
-func (s *FileChangeNode) IsLeaf() bool {
-	return s.File != nil
+func (s *FileChangeNode) ForEachFile(cb func(*models.File) error) error {
+	return forEachLeaf(s, func(n INode) error {
+		castNode := n.(*FileChangeNode)
+		return cb(castNode.File)
+	})
 }
 
-func (s *FileChangeNode) Size(collapsedPaths map[string]bool) int {
-	output := 1
-
-	if !collapsedPaths[s.GetPath()] {
-		for _, child := range s.Children {
-			output += child.Size(collapsedPaths)
-		}
-	}
-
-	return output
+func (s *FileChangeNode) Any(test func(node *FileChangeNode) bool) bool {
+	return any(s, func(n INode) bool {
+		castNode := n.(*FileChangeNode)
+		return test(castNode)
+	})
 }
 
-func (s *FileChangeNode) Flatten(collapsedPaths map[string]bool) []*FileChangeNode {
-	arr := []*FileChangeNode{s}
-
-	if !collapsedPaths[s.GetPath()] {
-		for _, child := range s.Children {
-			arr = append(arr, child.Flatten(collapsedPaths)...)
-		}
+func (n *FileChangeNode) Flatten(collapsedPaths map[string]bool) []*FileChangeNode {
+	results := flatten(n, collapsedPaths)
+	nodes := make([]*FileChangeNode, len(results))
+	for i, result := range results {
+		nodes[i] = result.(*FileChangeNode)
 	}
 
-	return arr
+	return nodes
 }
 
-func (s *FileChangeNode) Sort() {
-	s.sortChildren()
-
-	for _, child := range s.Children {
-		child.Sort()
-	}
+func (node *FileChangeNode) GetNodeAtIndex(index int, collapsedPaths map[string]bool) *FileChangeNode {
+	return getNodeAtIndex(node, index, collapsedPaths).(*FileChangeNode)
 }
 
-func (s *FileChangeNode) sortChildren() {
-	if s.IsLeaf() {
-		return
-	}
-
-	sortedChildren := make([]*FileChangeNode, len(s.Children))
-	copy(sortedChildren, s.Children)
-
-	sort.Slice(sortedChildren, func(i, j int) bool {
-		if !sortedChildren[i].IsLeaf() && sortedChildren[j].IsLeaf() {
-			return true
-		}
-		if sortedChildren[i].IsLeaf() && !sortedChildren[j].IsLeaf() {
-			return false
-		}
-
-		return sortedChildren[i].Path < sortedChildren[j].Path
-	})
-
-	// TODO: think about making this in-place
-	s.Children = sortedChildren
+func (node *FileChangeNode) GetIndexForPath(path string, collapsedPaths map[string]bool) (int, bool) {
+	return getIndexForPath(node, path, collapsedPaths)
 }
 
-func (s *FileChangeNode) GetPath() string {
-	return s.Path
+func (node *FileChangeNode) Size(collapsedPaths map[string]bool) int {
+	return size(node, collapsedPaths)
 }
 
 func (s *FileChangeNode) Compress() {
+	// with these functions I try to only have type conversion code on the actual struct,
+	// but comparing interface values to nil is fraught with danger so I'm duplicating
+	// that code here.
 	if s == nil {
 		return
 	}
 
-	s.compressAux()
+	compressAux(s)
 }
 
-func (s *FileChangeNode) compressAux() *FileChangeNode {
-	if s.IsLeaf() {
-		return s
-	}
-
-	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
-		}
-	}
+// This ignores the root
+func (node *FileChangeNode) GetPathsMatching(test func(*FileChangeNode) bool) []string {
+	return getPathsMatching(node, func(n INode) bool {
+		return test(n.(*FileChangeNode))
+	})
+}
 
-	for i := range s.Children {
-		s.Children[i] = s.Children[i].compressAux()
+func (s *FileChangeNode) GetLeaves() []*FileChangeNode {
+	leaves := getLeaves(s)
+	castLeaves := make([]*FileChangeNode, len(leaves))
+	for i := range leaves {
+		castLeaves[i] = leaves[i].(*FileChangeNode)
 	}
 
-	return s
-}
-
-func (s *FileChangeNode) HasExactlyOneChild() bool {
-	return len(s.Children) == 1
+	return castLeaves
 }
 
-// This ignores the root
-func (s *FileChangeNode) GetPathsMatching(test func(*FileChangeNode) bool) []string {
-	paths := []string{}
+// extra methods
 
-	if test(s) {
-		paths = append(paths, s.GetPath())
-	}
+func (s *FileChangeNode) GetHasUnstagedChanges() bool {
+	return s.AnyFile(func(file *models.File) bool { return file.HasUnstagedChanges })
+}
 
-	for _, child := range s.Children {
-		paths = append(paths, child.GetPathsMatching(test)...)
-	}
+func (s *FileChangeNode) GetHasStagedChanges() bool {
+	return s.AnyFile(func(file *models.File) bool { return file.HasStagedChanges })
+}
 
-	return paths
+func (s *FileChangeNode) GetHasInlineMergeConflicts() bool {
+	return s.AnyFile(func(file *models.File) bool { return file.HasInlineMergeConflicts })
 }
 
-func (s *FileChangeNode) ID() string {
-	return s.GetPath()
+func (s *FileChangeNode) GetIsTracked() bool {
+	return s.AnyFile(func(file *models.File) bool { return file.Tracked })
 }
 
-func (s *FileChangeNode) Description() string {
-	return s.GetPath()
+func (s *FileChangeNode) AnyFile(test func(file *models.File) bool) bool {
+	return s.Any(func(node *FileChangeNode) bool {
+		return node.IsLeaf() && test(node.File)
+	})
 }
 
-func (s *FileChangeNode) GetLeaves() []*FileChangeNode {
-	if s.IsLeaf() {
-		return []*FileChangeNode{s}
-	}
+func (s *FileChangeNode) NameAtDepth(depth int) string {
+	splitName := strings.Split(s.Path, string(os.PathSeparator))
+	name := filepath.Join(splitName[depth:]...)
 
-	output := []*FileChangeNode{}
-	for _, child := range s.Children {
-		output = append(output, child.GetLeaves()...)
+	if s.File != nil && s.File.IsRename() {
+		splitPrevName := strings.Split(s.File.PreviousName, string(os.PathSeparator))
+
+		prevName := s.File.PreviousName
+		// if the file has just been renamed inside the same directory, we can shave off
+		// the prefix for the previous path too. Otherwise we'll keep it unchanged
+		sameParentDir := len(splitName) == len(splitPrevName) && filepath.Join(splitName[0:depth]...) == filepath.Join(splitPrevName[0:depth]...)
+		if sameParentDir {
+			prevName = filepath.Join(splitPrevName[depth:]...)
+		}
+
+		return fmt.Sprintf("%s%s%s", prevName, " → ", name)
 	}
 
-	return output
+	return name
 }
diff --git a/pkg/gui/filetree/inode.go b/pkg/gui/filetree/inode.go
new file mode 100644
index 000000000..b3f9c73ee
--- /dev/null
+++ b/pkg/gui/filetree/inode.go
@@ -0,0 +1,196 @@
+package filetree
+
+import "sort"
+
+type INode interface {
+	IsLeaf() bool
+	GetPath() string
+	GetChildren() []INode
+	SetChildren([]INode)
+	GetCompressionLevel() int
+	SetCompressionLevel(int)
+}
+
+func sortNode(node INode) {
+	sortChildren(node)
+
+	for _, child := range node.GetChildren() {
+		sortNode(child)
+	}
+}
+
+func sortChildren(node INode) {
+	if node.IsLeaf() {
+		return
+	}
+
+	children := node.GetChildren()
+	sortedChildren := make([]INode, len(children))
+	copy(sortedChildren, children)
+
+	sort.Slice(sortedChildren, func(i, j int) bool {
+		if !sortedChildren[i].IsLeaf() && sortedChildren[j].IsLeaf() {
+			return true
+		}
+		if sortedChildren[i].IsLeaf() && !sortedChildren[j].IsLeaf() {
+			return false
+		}
+
+		return sortedChildren[i].GetPath() < sortedChildren[j].GetPath()
+	})
+
+	// TODO: think about making this in-place
+	node.SetChildren(sortedChildren)
+}
+
+func forEachLeaf(node INode, cb func(INode) error) error {
+	if node.IsLeaf() {
+		if err := cb(node); err != nil {
+			return err
+		}
+	}
+
+	for _, child := range node.GetChildren() {
+		if err := forEachLeaf(child, cb); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func any(node INode, test func(INode) bool) bool {
+	if test(node) {
+		return true
+	}
+
+	for _, child := range node.GetChildren() {
+		if any(child, test) {
+			return true
+		}
+	}
+
+	return false
+}
+
+func flatten(node INode, collapsedPaths map[string]bool) []INode {
+	result := []INode{}
+	result = append(result, node)
+
+	if !collapsedPaths[node.GetPath()] {
+		for _, child := range node.GetChildren() {
+			result = append(result, flatten(child, collapsedPaths)...)
+		}
+	}
+
+	return result
+}
+
+func getNodeAtIndex(node INode, index int, collapsedPaths map[string]bool) INode {
+	foundNode, _ := getNodeAtIndexAux(node, index, collapsedPaths)
+
+	return foundNode
+}
+
+func getNodeAtIndexAux(node INode, index int, collapsedPaths map[string]bool) (INode, int) {
+	offset := 1
+
+	if index == 0 {
+		return node, offset
+	}
+
+	if !collapsedPaths[node.GetPath()] {
+		for _, child := range node.GetChildren() {
+			foundNode, offsetChange := getNodeAtIndexAux(child, index-offset, collapsedPaths)
+			offset += offsetChange
+			if foundNode != nil {
+				return foundNode, offset
+			}
+		}
+	}
+
+	return nil, offset
+}
+
+func getIndexForPath(node INode, path string, collapsedPaths map[string]bool) (int, bool) {
+	offset := 0
+
+	if node.GetPath() == path {
+		return offset, true
+	}
+
+	if !collapsedPaths[node.GetPath()] {
+		for _, child := range node.GetChildren() {
+			offsetChange, found := getIndexForPath(child, path, collapsedPaths)
+			offset += offsetChange + 1
+			if found {
+				return offset, true
+			}
+		}
+	}
+
+	return offset, false
+}
+
+func size(node INode, collapsedPaths map[string]bool) int {
+	output := 1
+
+	if !collapsedPaths[node.GetPath()] {
+		for _, child := range node.GetChildren() {
+			output += size(child, collapsedPaths)
+		}
+	}
+
+	return output
+}
+
+func compressAux(node INode) INode {
+	if node.IsLeaf() {
+		return node
+	}
+
+	children := node.GetChildren()
+	for i := range children {
+		grandchildren := children[i].GetChildren()
+		for len(grandchildren) == 1 {
+			grandchildren[0].SetCompressionLevel(children[i].GetCompressionLevel() + 1)
+			children[i] = grandchildren[0]
+			grandchildren = children[i].GetChildren()
+		}
+	}
+
+	for i := range children {
+		children[i] = compressAux(children[i])
+	}
+
+	node.SetChildren(children)
+
+	return node
+}
+
+func getPathsMatching(node INode, test func(INode) bool) []string {
+	paths := []string{}
+
+	if test(node) {
+		paths = append(paths, node.GetPath())
+	}
+
+	for _, child := range node.GetChildren() {
+		paths = append(paths, getPathsMatching(child, test)...)
+	}
+
+	return paths
+}
+
+func getLeaves(node INode) []INode {
+	if node.IsLeaf() {
+		return []INode{node}
+	}
+
+	output := []INode{}
+	for _, child := range node.GetChildren() {
+		output = append(output, getLeaves(child)...)
+	}
+
+	return output
+}