1 files changed, 63 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}
+}