Maintain selection and simplify code

Several changes here:
- Simplify UncollapseAll; it's much too much work to iterate over files and call
  ExpandToPath for each one. We simply need to clear our container of collapsed
  paths.
- For CollapseAll, move the code that gets all dirPaths from the controller to
  the tree; I don't really see a reason why the controller should get the paths
  and pass them in
- Maintain the selection when collapsing or expanding all. For collapsing, we
  keep the top-level item selected that the previously selected node was a child
  of; for expanding, it's always possible to keep the item itself selected.

pkg/gui/controllers/commits_files_controller.go

@@ -416,27 +416,15 @@ func (self *CommitFilesController) toggleTreeView() error {
 }
 
 func (self *CommitFilesController) collapseAll() error {
-	nodes := self.context().GetAllItems()
-
-	dirPaths := lo.FilterMap(nodes, func(file *filetree.CommitFileNode, index int) (string, bool) {
-		return file.Path, !file.IsFile()
-	})
-
-	self.context().CommitFileTreeViewModel.CollapseAll(dirPaths)
+	self.context().CommitFileTreeViewModel.CollapseAll()
 
 	self.c.PostRefreshUpdate(self.context())
 
 	return nil
 }
 
 func (self *CommitFilesController) expandAll() error {
-	files := self.context().GetAllFiles()
-
-	filePaths := lo.Map(files, func(file *models.CommitFile, index int) string {
-		return file.GetPath()
-	})
-
-	self.context().CommitFileTreeViewModel.UncollapseAll(filePaths)
+	self.context().CommitFileTreeViewModel.UncollapseAll()
 
 	self.c.PostRefreshUpdate(self.context())
 

pkg/gui/controllers/files_controller.go

@@ -493,27 +493,15 @@ func (self *FilesController) enter() error {
 }
 
 func (self *FilesController) collapseAll() error {
-	nodes := self.context().GetAllItems()
-
-	dirPaths := lo.FilterMap(nodes, func(file *filetree.FileNode, index int) (string, bool) {
-		return file.Path, !file.IsFile()
-	})
-
-	self.context().FileTreeViewModel.CollapseAll(dirPaths)
+	self.context().FileTreeViewModel.CollapseAll()
 
 	self.c.PostRefreshUpdate(self.context())
 
 	return nil
 }
 
 func (self *FilesController) expandAll() error {
-	files := self.context().GetAllFiles()
-
-	filePaths := lo.Map(files, func(file *models.File, index int) string {
-		return file.GetPath()
-	})
-
-	self.context().FileTreeViewModel.UncollapseAll(filePaths)
+	self.context().FileTreeViewModel.UncollapseAll()
 
 	self.c.PostRefreshUpdate(self.context())
 

pkg/gui/filetree/collapsed_paths.go

@@ -36,3 +36,8 @@ func (self *CollapsedPaths) ToggleCollapsed(path string) {
 		self.collapsedPaths.Add(path)
 	}
 }
+
+func (self *CollapsedPaths) ExpandAll() {
+	// Could be cleaner if Set had a Clear() method...
+	self.collapsedPaths.RemoveSlice(self.collapsedPaths.ToSlice())
+}

pkg/gui/filetree/commit_file_tree.go

@@ -25,16 +25,18 @@ type CommitFileTree struct {
 	collapsedPaths *CollapsedPaths
 }
 
-func (self *CommitFileTree) CollapseAll(paths []string) {
-	for _, path := range paths {
+func (self *CommitFileTree) CollapseAll() {
+	dirPaths := lo.FilterMap(self.GetAllItems(), func(file *CommitFileNode, index int) (string, bool) {
+		return file.Path, !file.IsFile()
+	})
+
+	for _, path := range dirPaths {
 		self.collapsedPaths.Collapse(path)
 	}
 }
 
-func (self *CommitFileTree) UncollapseAll(paths []string) {
-	for _, path := range paths {
-		self.collapsedPaths.ExpandToPath(path)
-	}
+func (self *CommitFileTree) UncollapseAll() {
+	self.collapsedPaths.ExpandAll()
 }
 
 var _ ICommitFileTree = &CommitFileTree{}

pkg/gui/filetree/commit_file_tree_view_model.go

@@ -1,6 +1,7 @@
 package filetree
 
 import (
+	"strings"
 	"sync"
 
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
@@ -160,3 +161,33 @@ func (self *CommitFileTreeViewModel) ToggleShowTree() {
 		self.SetSelection(index)
 	}
 }
+
+func (self *CommitFileTreeViewModel) CollapseAll() {
+	selectedNode := self.GetSelected()
+
+	self.ICommitFileTree.CollapseAll()
+	if selectedNode == nil {
+		return
+	}
+
+	topLevelPath := strings.Split(selectedNode.Path, "/")[0]
+	index, found := self.GetIndexForPath(topLevelPath)
+	if found {
+		self.SetSelectedLineIdx(index)
+	}
+}
+
+func (self *CommitFileTreeViewModel) UncollapseAll() {
+	selectedNode := self.GetSelected()
+
+	self.ICommitFileTree.UncollapseAll()
+
+	if selectedNode == nil {
+		return
+	}
+
+	index, found := self.GetIndexForPath(selectedNode.Path)
+	if found {
+		self.SetSelectedLineIdx(index)
+	}
+}

pkg/gui/filetree/file_tree.go

@@ -31,8 +31,8 @@ type ITree[T any] interface {
 	IsCollapsed(path string) bool
 	ToggleCollapsed(path string)
 	CollapsedPaths() *CollapsedPaths
-	CollapseAll(paths []string)
-	UncollapseAll(paths []string)
+	CollapseAll()
+	UncollapseAll()
 }
 
 type IFileTree interface {
@@ -173,16 +173,18 @@ func (self *FileTree) ToggleCollapsed(path string) {
 	self.collapsedPaths.ToggleCollapsed(path)
 }
 
-func (self *FileTree) CollapseAll(paths []string) {
-	for _, path := range paths {
+func (self *FileTree) CollapseAll() {
+	dirPaths := lo.FilterMap(self.GetAllItems(), func(file *FileNode, index int) (string, bool) {
+		return file.Path, !file.IsFile()
+	})
+
+	for _, path := range dirPaths {
 		self.collapsedPaths.Collapse(path)
 	}
 }
 
-func (self *FileTree) UncollapseAll(paths []string) {
-	for _, path := range paths {
-		self.collapsedPaths.ExpandToPath(path)
-	}
+func (self *FileTree) UncollapseAll() {
+	self.collapsedPaths.ExpandAll()
 }
 
 func (self *FileTree) Tree() *FileNode {

pkg/gui/filetree/file_tree_view_model.go

@@ -1,6 +1,7 @@
 package filetree
 
 import (
+	"strings"
 	"sync"
 
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
@@ -190,3 +191,33 @@ func (self *FileTreeViewModel) ToggleShowTree() {
 		self.SetSelectedLineIdx(index)
 	}
 }
+
+func (self *FileTreeViewModel) CollapseAll() {
+	selectedNode := self.GetSelected()
+
+	self.IFileTree.CollapseAll()
+	if selectedNode == nil {
+		return
+	}
+
+	topLevelPath := strings.Split(selectedNode.Path, "/")[0]
+	index, found := self.GetIndexForPath(topLevelPath)
+	if found {
+		self.SetSelectedLineIdx(index)
+	}
+}
+
+func (self *FileTreeViewModel) UncollapseAll() {
+	selectedNode := self.GetSelected()
+
+	self.IFileTree.UncollapseAll()
+
+	if selectedNode == nil {
+		return
+	}
+
+	index, found := self.GetIndexForPath(selectedNode.Path)
+	if found {
+		self.SetSelectedLineIdx(index)
+	}
+}