small refactor

3 files changed, 292 insertions, 0 deletions

diff --git a/pkg/gui/filetree/build_tree.go b/pkg/gui/filetree/build_tree.go
new file mode 100644
index 000000000..e8da844e0
--- /dev/null
+++ b/pkg/gui/filetree/build_tree.go
@@ -0,0 +1,63 @@
+package filetree
+
+import (
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+)
+
+func BuildTreeFromFiles(files []*models.File) *models.FileChangeNode {
+	root := &models.FileChangeNode{}
+
+	var curr *models.FileChangeNode
+	for _, file := range files {
+		split := strings.Split(file.Name, string(os.PathSeparator))
+		curr = root
+	outer:
+		for i := range split {
+			var setFile *models.File
+			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 := &models.FileChangeNode{
+				Path: path,
+				File: setFile,
+			}
+			curr.Children = append(curr.Children, newChild)
+
+			curr = newChild
+		}
+	}
+
+	root.Sort()
+	root.Compress()
+
+	return root
+}
+
+func BuildFlatTreeFromFiles(files []*models.File) *models.FileChangeNode {
+	rootAux := BuildTreeFromFiles(files)
+	sortedFiles := rootAux.GetLeaves()
+
+	// Move merge conflicts to top. This is the one way in which sorting
+	// differs between flat mode and tree mode
+	sort.SliceStable(sortedFiles, func(i, j int) bool {
+		return sortedFiles[i].File != nil && sortedFiles[i].File.HasMergeConflicts && !(sortedFiles[j].File != nil && sortedFiles[j].File.HasMergeConflicts)
+	})
+
+	return &models.FileChangeNode{Children: sortedFiles}
+}
diff --git a/pkg/gui/filetree/file_change_manager.go b/pkg/gui/filetree/file_change_manager.go
new file mode 100644
index 000000000..925398827
--- /dev/null
+++ b/pkg/gui/filetree/file_change_manager.go
@@ -0,0 +1,138 @@
+package filetree
+
+import (
+	"fmt"
+	"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 = "►"
+
+type FileChangeManager struct {
+	Files          []*models.File
+	Tree           *models.FileChangeNode
+	ShowTree       bool
+	Log            *logrus.Entry
+	CollapsedPaths map[string]bool
+}
+
+func NewFileChangeManager(files []*models.File, log *logrus.Entry, showTree bool) *FileChangeManager {
+	return &FileChangeManager{
+		Files:          files,
+		Log:            log,
+		ShowTree:       showTree,
+		CollapsedPaths: map[string]bool{},
+	}
+}
+
+func (m *FileChangeManager) GetItemAtIndex(index int) *models.FileChangeNode {
+	// 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 *FileChangeManager) GetIndexForPath(path string) (int, bool) {
+	index, found := m.Tree.GetIndexForPath(path, m.CollapsedPaths)
+	return index - 1, found
+}
+
+func (m *FileChangeManager) GetAllItems() []*models.FileChangeNode {
+	if m.Tree == nil {
+		return nil
+	}
+
+	return m.Tree.Flatten(m.CollapsedPaths)[1:] // ignoring root
+}
+
+func (m *FileChangeManager) GetItemsLength() int {
+	return m.Tree.Size(m.CollapsedPaths) - 1 // ignoring root
+}
+
+func (m *FileChangeManager) GetAllFiles() []*models.File {
+	return m.Files
+}
+
+func (m *FileChangeManager) SetFiles(files []*models.File) {
+	m.Files = files
+
+	m.SetTree()
+}
+
+func (m *FileChangeManager) SetTree() {
+	if m.ShowTree {
+		m.Tree = BuildTreeFromFiles(m.Files)
+	} else {
+		m.Tree = BuildFlatTreeFromFiles(m.Files)
+	}
+}
+
+func (m *FileChangeManager) Render(diffName string, submoduleConfigs []*models.SubmoduleConfig) []string {
+	return m.renderAux(m.Tree, "", -1, diffName, submoduleConfigs)
+}
+
+const INNER_ITEM = "├─ "
+const LAST_ITEM = "└─ "
+const NESTED = "│  "
+const NOTHING = "   "
+
+func (m *FileChangeManager) IsCollapsed(s *models.FileChangeNode) bool {
+	return m.CollapsedPaths[s.GetPath()]
+}
+
+func (m *FileChangeManager) ToggleCollapsed(s *models.FileChangeNode) {
+	m.CollapsedPaths[s.GetPath()] = !m.CollapsedPaths[s.GetPath()]
+}
+
+func (m *FileChangeManager) renderAux(s *models.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
+}
diff --git a/pkg/gui/filetree/file_change_manager_test.go b/pkg/gui/filetree/file_change_manager_test.go
new file mode 100644
index 000000000..8e8ba00cd
--- /dev/null
+++ b/pkg/gui/filetree/file_change_manager_test.go
@@ -0,0 +1,91 @@
+package filetree
+
+import (
+	"testing"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRender(t *testing.T) {
+	scenarios := []struct {
+		name           string
+		root           *models.FileChangeNode
+		collapsedPaths map[string]bool
+		expected       []string
+	}{
+		{
+			name:     "nil node",
+			root:     nil,
+			expected: []string{},
+		},
+		{
+			name: "leaf node",
+			root: &models.FileChangeNode{
+				Path: "",
+				Children: []*models.FileChangeNode{
+					{File: &models.File{Name: "test", ShortStatus: " M", HasStagedChanges: true}, Path: "test"},
+				},
+			},
+			expected: []string{" M test"},
+		},
+		{
+			name: "big example",
+			root: &models.FileChangeNode{
+				Path: "",
+				Children: []*models.FileChangeNode{
+					{
+						Path: "dir1",
+						Children: []*models.FileChangeNode{
+							{
+								File: &models.File{Name: "dir1/file2", ShortStatus: "M ", HasUnstagedChanges: true},
+								Path: "dir1/file2",
+							},
+							{
+								File: &models.File{Name: "dir1/file3", ShortStatus: "M ", HasUnstagedChanges: true},
+								Path: "dir1/file3",
+							},
+						},
+					},
+					{
+						Path: "dir2",
+						Children: []*models.FileChangeNode{
+							{
+								Path: "dir2/dir2",
+								Children: []*models.FileChangeNode{
+									{
+										File: &models.File{Name: "dir2/dir2/file3", ShortStatus: " M", HasStagedChanges: true},
+										Path: "dir2/dir2/file3",
+									},
+									{
+										File: &models.File{Name: "dir2/dir2/file4", ShortStatus: "M ", HasUnstagedChanges: true},
+										Path: "dir2/dir2/file4",
+									},
+								},
+							},
+							{
+								File: &models.File{Name: "dir2/file5", ShortStatus: "M ", HasUnstagedChanges: true},
+								Path: "dir2/file5",
+							},
+						},
+					},
+					{
+						File: &models.File{Name: "file1", ShortStatus: "M ", HasUnstagedChanges: true},
+						Path: "file1",
+					},
+				},
+			},
+			expected:       []string{"dir1 ►", "dir2 ▼", "├─ dir2 ▼", "│  ├─  M file3", "│  └─ M  file4", "└─ M  file5", "M  file1"},
+			collapsedPaths: map[string]bool{"dir1": true},
+		},
+	}
+
+	for _, s := range scenarios {
+		s := s
+		t.Run(s.name, func(t *testing.T) {
+			mngr := &FileChangeManager{Tree: s.root, CollapsedPaths: s.collapsedPaths}
+			result := mngr.Render("", nil)
+			assert.EqualValues(t, s.expected, result)
+		})
+	}
+}