file tree for commit files

6 files changed, 247 insertions, 81 deletions

diff --git a/pkg/gui/filetree/build_tree.go b/pkg/gui/filetree/build_tree.go
index 267bb22dd..d6fe9b98d 100644
--- a/pkg/gui/filetree/build_tree.go
+++ b/pkg/gui/filetree/build_tree.go
@@ -49,6 +49,53 @@ func BuildTreeFromFiles(files []*models.File) *FileChangeNode {
 	return root
 }
 
+func BuildFlatTreeFromCommitFiles(files []*models.CommitFile) *CommitFileChangeNode {
+	rootAux := BuildTreeFromCommitFiles(files)
+	sortedFiles := rootAux.GetLeaves()
+
+	return &CommitFileChangeNode{Children: sortedFiles}
+}
+
+func BuildTreeFromCommitFiles(files []*models.CommitFile) *CommitFileChangeNode {
+	root := &CommitFileChangeNode{}
+
+	var curr *CommitFileChangeNode
+	for _, file := range files {
+		split := strings.Split(file.Name, string(os.PathSeparator))
+		curr = root
+	outer:
+		for i := range split {
+			var setFile *models.CommitFile
+			isFile := i == len(split)-1
+			if isFile {
+				setFile = file
+			}
+
+			path := filepath.Join(split[:i+1]...)
+
+			for _, existingChild := range curr.Children {
+				if existingChild.Path == path {
+					curr = existingChild
+					continue outer
+				}
+			}
+
+			newChild := &CommitFileChangeNode{
+				Path: path,
+				File: setFile,
+			}
+			curr.Children = append(curr.Children, newChild)
+
+			curr = newChild
+		}
+	}
+
+	root.Sort()
+	root.Compress()
+
+	return root
+}
+
 func BuildFlatTreeFromFiles(files []*models.File) *FileChangeNode {
 	rootAux := BuildTreeFromFiles(files)
 	sortedFiles := rootAux.GetLeaves()
diff --git a/pkg/gui/filetree/collapsed_paths.go b/pkg/gui/filetree/collapsed_paths.go
new file mode 100644
index 000000000..60860414f
--- /dev/null
+++ b/pkg/gui/filetree/collapsed_paths.go
@@ -0,0 +1,25 @@
+package filetree
+
+import (
+	"os"
+	"strings"
+)
+
+type CollapsedPaths map[string]bool
+
+func (cp CollapsedPaths) ExpandToPath(path string) {
+	// need every directory along the way
+	split := strings.Split(path, string(os.PathSeparator))
+	for i := range split {
+		dir := strings.Join(split[0:i+1], string(os.PathSeparator))
+		cp[dir] = false
+	}
+}
+
+func (cp CollapsedPaths) IsCollapsed(path string) bool {
+	return cp[path]
+}
+
+func (cp CollapsedPaths) ToggleCollapsed(path string) {
+	cp[path] = !cp[path]
+}
diff --git a/pkg/gui/filetree/commit_file_change_manager.go b/pkg/gui/filetree/commit_file_change_manager.go
new file mode 100644
index 000000000..d6b70a2bd
--- /dev/null
+++ b/pkg/gui/filetree/commit_file_change_manager.go
@@ -0,0 +1,95 @@
+package filetree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/jesseduffield/lazygit/pkg/gui/presentation"
+	"github.com/sirupsen/logrus"
+)
+
+type CommitFileChangeManager struct {
+	files          []*models.CommitFile
+	tree           *CommitFileChangeNode
+	showTree       bool
+	log            *logrus.Entry
+	collapsedPaths CollapsedPaths
+	// parent is the identifier of the parent object e.g. a commit SHA if this commit file is for a commit, or a stash entry ref like 'stash@{1}'
+	parent string
+}
+
+func (m *CommitFileChangeManager) GetParent() string {
+	return m.parent
+}
+
+func NewCommitFileChangeManager(files []*models.CommitFile, log *logrus.Entry, showTree bool) *CommitFileChangeManager {
+	return &CommitFileChangeManager{
+		files:          files,
+		log:            log,
+		showTree:       showTree,
+		collapsedPaths: CollapsedPaths{},
+	}
+}
+
+func (m *CommitFileChangeManager) ExpandToPath(path string) {
+	m.collapsedPaths.ExpandToPath(path)
+}
+
+func (m *CommitFileChangeManager) ToggleShowTree() {
+	m.showTree = !m.showTree
+	m.SetTree()
+}
+
+func (m *CommitFileChangeManager) GetItemAtIndex(index int) *CommitFileChangeNode {
+	// need to traverse the three depth first until we get to the index.
+	return m.tree.GetNodeAtIndex(index+1, m.collapsedPaths) // ignoring root
+}
+
+func (m *CommitFileChangeManager) GetIndexForPath(path string) (int, bool) {
+	index, found := m.tree.GetIndexForPath(path, m.collapsedPaths)
+	return index - 1, found
+}
+
+func (m *CommitFileChangeManager) GetAllItems() []*CommitFileChangeNode {
+	if m.tree == nil {
+		return nil
+	}
+
+	return m.tree.Flatten(m.collapsedPaths)[1:] // ignoring root
+}
+
+func (m *CommitFileChangeManager) GetItemsLength() int {
+	return m.tree.Size(m.collapsedPaths) - 1 // ignoring root
+}
+
+func (m *CommitFileChangeManager) GetAllFiles() []*models.CommitFile {
+	return m.files
+}
+
+func (m *CommitFileChangeManager) SetFiles(files []*models.CommitFile, parent string) {
+	m.files = files
+	m.parent = parent
+
+	m.SetTree()
+}
+
+func (m *CommitFileChangeManager) SetTree() {
+	if m.showTree {
+		m.tree = BuildTreeFromCommitFiles(m.files)
+	} else {
+		m.tree = BuildFlatTreeFromCommitFiles(m.files)
+	}
+}
+
+func (m *CommitFileChangeManager) IsCollapsed(path string) bool {
+	return m.collapsedPaths.IsCollapsed(path)
+}
+
+func (m *CommitFileChangeManager) ToggleCollapsed(path string) {
+	m.collapsedPaths.ToggleCollapsed(path)
+}
+
+func (m *CommitFileChangeManager) Render(diffName string) []string {
+	return renderAux(m.tree, m.collapsedPaths, "", -1, func(n INode, depth int) string {
+		castN := n.(*CommitFileChangeNode)
+		return presentation.GetCommitFileLine(castN.NameAtDepth(depth), diffName, castN.File)
+	})
+}
diff --git a/pkg/gui/filetree/constants.go b/pkg/gui/filetree/constants.go
new file mode 100644
index 000000000..d510650e2
--- /dev/null
+++ b/pkg/gui/filetree/constants.go
@@ -0,0 +1,9 @@
+package filetree
+
+const EXPANDED_ARROW = "▼"
+const COLLAPSED_ARROW = "►"
+
+const INNER_ITEM = "├─ "
+const LAST_ITEM = "└─ "
+const NESTED = "│  "
+const NOTHING = "   "
diff --git a/pkg/gui/filetree/file_change_manager.go b/pkg/gui/filetree/file_change_manager.go
index 30bd7cf1e..b663e127b 100644
--- a/pkg/gui/filetree/file_change_manager.go
+++ b/pkg/gui/filetree/file_change_manager.go
@@ -1,29 +1,17 @@
 package filetree
 
 import (
-	"fmt"
-	"os"
-	"strings"
-
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation"
 	"github.com/sirupsen/logrus"
 )
 
-const EXPANDED_ARROW = "▼"
-const COLLAPSED_ARROW = "►"
-
-const INNER_ITEM = "├─ "
-const LAST_ITEM = "└─ "
-const NESTED = "│  "
-const NOTHING = "   "
-
 type FileChangeManager struct {
 	files          []*models.File
 	tree           *FileChangeNode
 	showTree       bool
 	log            *logrus.Entry
-	collapsedPaths map[string]bool
+	collapsedPaths CollapsedPaths
 }
 
 func NewFileChangeManager(files []*models.File, log *logrus.Entry, showTree bool) *FileChangeManager {
@@ -31,10 +19,14 @@ func NewFileChangeManager(files []*models.File, log *logrus.Entry, showTree bool
 		files:          files,
 		log:            log,
 		showTree:       showTree,
-		collapsedPaths: map[string]bool{},
+		collapsedPaths: CollapsedPaths{},
 	}
 }
 
+func (m *FileChangeManager) ExpandToPath(path string) {
+	m.collapsedPaths.ExpandToPath(path)
+}
+
 func (m *FileChangeManager) ToggleShowTree() {
 	m.showTree = !m.showTree
 	m.SetTree()
@@ -80,74 +72,17 @@ func (m *FileChangeManager) SetTree() {
 	}
 }
 
-func (m *FileChangeManager) Render(diffName string, submoduleConfigs []*models.SubmoduleConfig) []string {
-	return m.renderAux(m.tree, "", -1, diffName, submoduleConfigs)
+func (m *FileChangeManager) IsCollapsed(path string) bool {
+	return m.collapsedPaths.IsCollapsed(path)
 }
 
-func (m *FileChangeManager) IsCollapsed(s *FileChangeNode) bool {
-	return m.collapsedPaths[s.GetPath()]
-}
-
-func (m *FileChangeManager) ToggleCollapsed(s *FileChangeNode) {
-	m.collapsedPaths[s.GetPath()] = !m.collapsedPaths[s.GetPath()]
-}
-
-func (m *FileChangeManager) renderAux(s *FileChangeNode, prefix string, depth int, diffName string, submoduleConfigs []*models.SubmoduleConfig) []string {
-	isRoot := depth == -1
-	if s == nil {
-		return []string{}
-	}
-
-	getLine := func() string {
-		return prefix + presentation.GetFileLine(s.GetHasUnstagedChanges(), s.GetHasStagedChanges(), s.NameAtDepth(depth), diffName, submoduleConfigs, s.File)
-	}
-
-	if s.IsLeaf() {
-		if isRoot {
-			return []string{}
-		}
-		return []string{getLine()}
-	}
-
-	if m.IsCollapsed(s) {
-		return []string{fmt.Sprintf("%s %s", getLine(), COLLAPSED_ARROW)}
-	}
-
-	arr := []string{}
-	if !isRoot {
-		arr = append(arr, fmt.Sprintf("%s %s", getLine(), EXPANDED_ARROW))
-	}
-
-	newPrefix := prefix
-	if strings.HasSuffix(prefix, LAST_ITEM) {
-		newPrefix = strings.TrimSuffix(prefix, LAST_ITEM) + NOTHING
-	} else if strings.HasSuffix(prefix, INNER_ITEM) {
-		newPrefix = strings.TrimSuffix(prefix, INNER_ITEM) + NESTED
-	}
-
-	for i, child := range s.Children {
-		isLast := i == len(s.Children)-1
-
-		var childPrefix string
-		if isRoot {
-			childPrefix = newPrefix
-		} else if isLast {
-			childPrefix = newPrefix + LAST_ITEM
-		} else {
-			childPrefix = newPrefix + INNER_ITEM
-		}
-
-		arr = append(arr, m.renderAux(child, childPrefix, depth+1+s.CompressionLevel, diffName, submoduleConfigs)...)
-	}
-
-	return arr
+func (m *FileChangeManager) ToggleCollapsed(path string) {
+	m.collapsedPaths.ToggleCollapsed(path)
 }
 
-func (m *FileChangeManager) ExpandToPath(path string) {
-	// need every directory along the way
-	split := strings.Split(path, string(os.PathSeparator))
-	for i := range split {
-		dir := strings.Join(split[0:i+1], string(os.PathSeparator))
-		m.collapsedPaths[dir] = false
-	}
+func (m *FileChangeManager) Render(diffName string, submoduleConfigs []*models.SubmoduleConfig) []string {
+	return renderAux(m.tree, m.collapsedPaths, "", -1, func(n INode, depth int) string {
+		castN := n.(*FileChangeNode)
+		return presentation.GetFileLine(castN.GetHasUnstagedChanges(), castN.GetHasStagedChanges(), castN.NameAtDepth(depth), diffName, submoduleConfigs, castN.File)
+	})
 }
diff --git a/pkg/gui/filetree/inode.go b/pkg/gui/filetree/inode.go
index b3f9c73ee..bbd1c3fb6 100644
--- a/pkg/gui/filetree/inode.go
+++ b/pkg/gui/filetree/inode.go
@@ -1,6 +1,10 @@
 package filetree
 
-import "sort"
+import (
+	"fmt"
+	"sort"
+	"strings"
+)
 
 type INode interface {
 	IsLeaf() bool
@@ -194,3 +198,54 @@ func getLeaves(node INode) []INode {
 
 	return output
 }
+
+func renderAux(s INode, collapsedPaths CollapsedPaths, prefix string, depth int, renderLine func(INode, int) string) []string {
+	isRoot := depth == -1
+	if s == nil {
+		return []string{}
+	}
+
+	renderLineWithPrefix := func() string {
+		return prefix + renderLine(s, depth)
+	}
+
+	if s.IsLeaf() {
+		if isRoot {
+			return []string{}
+		}
+		return []string{renderLineWithPrefix()}
+	}
+
+	if collapsedPaths.IsCollapsed(s.GetPath()) {
+		return []string{fmt.Sprintf("%s %s", renderLineWithPrefix(), COLLAPSED_ARROW)}
+	}
+
+	arr := []string{}
+	if !isRoot {
+		arr = append(arr, fmt.Sprintf("%s %s", renderLineWithPrefix(), EXPANDED_ARROW))
+	}
+
+	newPrefix := prefix
+	if strings.HasSuffix(prefix, LAST_ITEM) {
+		newPrefix = strings.TrimSuffix(prefix, LAST_ITEM) + NOTHING
+	} else if strings.HasSuffix(prefix, INNER_ITEM) {
+		newPrefix = strings.TrimSuffix(prefix, INNER_ITEM) + NESTED
+	}
+
+	for i, child := range s.GetChildren() {
+		isLast := i == len(s.GetChildren())-1
+
+		var childPrefix string
+		if isRoot {
+			childPrefix = newPrefix
+		} else if isLast {
+			childPrefix = newPrefix + LAST_ITEM
+		} else {
+			childPrefix = newPrefix + INNER_ITEM
+		}
+
+		arr = append(arr, renderAux(child, collapsedPaths, childPrefix, depth+1+s.GetCompressionLevel(), renderLine)...)
+	}
+
+	return arr
+}