Improve performance for discarding range selection of many files

This is probably the less severe case, but it could still be an issue for people
who have many modified top-level files they want to discard, and have
showRootItemInFileTree set to false; they could select all those files by
pressing 'v' and '>'.

pkg/commands/git_commands/working_tree.go

@@ -184,47 +184,49 @@ type IFileNode interface {
 	GetFile() *models.File
 }
 
-func (self *WorkingTreeCommands) DiscardAllDirChanges(node IFileNode) error {
+func (self *WorkingTreeCommands) DiscardAllDirChanges(nodes []IFileNode) error {
 	// Collect files into buckets so we can batch git calls where possible.
 	var specialFiles []*models.File // renames, AA, DU — handled individually
 	var filesToReset []string       // need `git reset` first (staged or conflicted)
 	var filesToCheckout []string    // need `git checkout` (after optional reset)
 	var filesToRemove []string      // added files to delete from disk
 
-	_ = node.ForEachFile(func(file *models.File) error {
-		// Renames and certain merge-conflict statuses need per-file logic.
-		if file.IsRename() || file.ShortStatus == "AA" || file.ShortStatus == "DU" {
-			specialFiles = append(specialFiles, file)
-			return nil
-		}
+	for _, node := range nodes {
+		_ = node.ForEachFile(func(file *models.File) error {
+			// Renames and certain merge-conflict statuses need per-file logic.
+			if file.IsRename() || file.ShortStatus == "AA" || file.ShortStatus == "DU" {
+				specialFiles = append(specialFiles, file)
+				return nil
+			}
 
-		if file.HasStagedChanges || file.HasMergeConflicts {
-			filesToReset = append(filesToReset, file.Path)
-			// DD and AU are done after the reset; no checkout or remove needed.
+			if file.HasStagedChanges || file.HasMergeConflicts {
+				filesToReset = append(filesToReset, file.Path)
+				// DD and AU are done after the reset; no checkout or remove needed.
+				if file.ShortStatus == "DD" || file.ShortStatus == "AU" {
+					return nil
+				}
+				if file.Added {
+					filesToRemove = append(filesToRemove, file.Path)
+				} else {
+					filesToCheckout = append(filesToCheckout, file.Path)
+				}
+				return nil
+			}
+
+			// No staged changes below this point.
 			if file.ShortStatus == "DD" || file.ShortStatus == "AU" {
 				return nil
 			}
+
 			if file.Added {
 				filesToRemove = append(filesToRemove, file.Path)
-			} else {
-				filesToCheckout = append(filesToCheckout, file.Path)
+				return nil
 			}
-			return nil
-		}
 
-		// No staged changes below this point.
-		if file.ShortStatus == "DD" || file.ShortStatus == "AU" {
+			filesToCheckout = append(filesToCheckout, file.Path)
 			return nil
-		}
-
-		if file.Added {
-			filesToRemove = append(filesToRemove, file.Path)
-			return nil
-		}
-
-		filesToCheckout = append(filesToCheckout, file.Path)
-		return nil
-	})
+		})
+	}
 
 	for _, file := range specialFiles {
 		if err := self.DiscardAllFileChanges(file); err != nil {
@@ -245,34 +247,34 @@ func (self *WorkingTreeCommands) DiscardAllDirChanges(node IFileNode) error {
 	return runGitCmdOnPaths("checkout", filesToCheckout, self.cmd)
 }
 
-func (self *WorkingTreeCommands) DiscardUnstagedDirChanges(node IFileNode) error {
-	file := node.GetFile()
-	if file == nil {
-		if err := self.RemoveUntrackedDirFiles(node); err != nil {
-			return err
-		}
+func (self *WorkingTreeCommands) DiscardUnstagedDirChanges(nodes []IFileNode) error {
+	// Collect files into buckets so we can batch git calls where possible.
+	// Use specific file paths rather than directory paths, so that an active
+	// filter (e.g. from pressing `/`) only discards visible files.
+	var filesToRemove []string   // purely untracked: remove from disk
+	var filesToCheckout []string // tracked or staged: restore via checkout
 
-		// Use specific file paths rather than the directory path, so that an
-		// active filter (e.g. from pressing `/`) only discards visible files.
-		// Include staged files: a file that is staged but also has additional
-		// unstaged changes (AM status) needs checkout to discard those changes.
-		trackedPaths := node.GetFilePathsMatching(func(f *models.File) bool {
-			return f.GetIsTracked() || f.GetHasStagedChanges()
+	for _, node := range nodes {
+		_ = node.ForEachFile(func(file *models.File) error {
+			if !file.Tracked && !file.HasStagedChanges {
+				filesToRemove = append(filesToRemove, file.Path)
+			} else {
+				// Include staged files: a file that is staged but also has
+				// additional unstaged changes (AM status) needs checkout to
+				// discard those changes.
+				filesToCheckout = append(filesToCheckout, file.Path)
+			}
+			return nil
 		})
-		if err := runGitCmdOnPaths("checkout", trackedPaths, self.cmd); err != nil {
-			return err
-		}
-	} else {
-		if file.Added && !file.HasStagedChanges {
-			return self.os.RemoveFile(file.Path)
-		}
+	}
 
-		if err := self.DiscardUnstagedFileChanges(file); err != nil {
+	for _, path := range filesToRemove {
+		if err := self.os.RemoveFile(path); err != nil {
 			return err
 		}
 	}
 
-	return nil
+	return runGitCmdOnPaths("checkout", filesToCheckout, self.cmd)
 }
 
 func (self *WorkingTreeCommands) RemoveUntrackedDirFiles(node IFileNode) error {

pkg/commands/git_commands/working_tree_test.go

@@ -501,7 +501,7 @@ func (n *testNode) GetFile() *models.File { return n.file }
 func TestWorkingTreeDiscardAllDirChanges(t *testing.T) {
 	type scenario struct {
 		testName             string
-		node                 *testNode
+		nodes                []IFileNode
 		runner               *oscommands.FakeCmdObjRunner
 		expectedRemovedFiles []string
 	}
@@ -509,39 +509,61 @@ func TestWorkingTreeDiscardAllDirChanges(t *testing.T) {
 	scenarios := []scenario{
 		{
 			testName: "multiple regular tracked files batched into a single checkout call",
-			node: &testNode{
+			nodes: []IFileNode{&testNode{
 				files: []*models.File{
 					{Path: "a.txt", Tracked: true},
 					{Path: "b.txt", Tracked: true},
 					{Path: "c.txt", Tracked: true},
 				},
-			},
+			}},
 			runner: oscommands.NewFakeRunner(t).
 				ExpectGitArgs([]string{"checkout", "--", "a.txt", "b.txt", "c.txt"}, "", nil),
 		},
 		{
 			testName: "staged files batched into a single reset then a single checkout",
-			node: &testNode{
+			nodes: []IFileNode{&testNode{
 				files: []*models.File{
 					{Path: "a.txt", Tracked: true, HasStagedChanges: true},
 					{Path: "b.txt", Tracked: true, HasStagedChanges: true},
 				},
-			},
+			}},
 			runner: oscommands.NewFakeRunner(t).
 				ExpectGitArgs([]string{"reset", "--", "a.txt", "b.txt"}, "", nil).
 				ExpectGitArgs([]string{"checkout", "--", "a.txt", "b.txt"}, "", nil),
 		},
 		{
 			testName: "added files with no staged changes are removed from disk without any git call",
-			node: &testNode{
+			nodes: []IFileNode{&testNode{
 				files: []*models.File{
 					{Path: "new1.txt", Added: true},
 					{Path: "new2.txt", Added: true},
 				},
-			},
+			}},
 			runner:               oscommands.NewFakeRunner(t),
 			expectedRemovedFiles: []string{"new1.txt", "new2.txt"},
 		},
+		{
+			testName: "files from multiple nodes are batched into a single git call",
+			nodes: []IFileNode{
+				&testNode{
+					path: "dir1",
+					files: []*models.File{
+						{Path: "dir1/a.txt", Tracked: true},
+						{Path: "dir1/b.txt", Added: true},
+					},
+				},
+				&testNode{
+					path: "dir2",
+					files: []*models.File{
+						{Path: "dir2/c.txt", Tracked: true},
+						{Path: "dir2/d.txt", Added: true},
+					},
+				},
+			},
+			runner: oscommands.NewFakeRunner(t).
+				ExpectGitArgs([]string{"checkout", "--", "dir1/a.txt", "dir2/c.txt"}, "", nil),
+			expectedRemovedFiles: []string{"dir1/b.txt", "dir2/d.txt"},
+		},
 	}
 
 	for _, s := range scenarios {
@@ -552,7 +574,7 @@ func TestWorkingTreeDiscardAllDirChanges(t *testing.T) {
 				return nil
 			}
 			instance := buildWorkingTreeCommands(commonDeps{runner: s.runner, removeFile: removeFile})
-			err := instance.DiscardAllDirChanges(s.node)
+			err := instance.DiscardAllDirChanges(s.nodes)
 			assert.NoError(t, err)
 			assert.Equal(t, s.expectedRemovedFiles, removedFiles)
 			s.runner.CheckForMissingCalls()
@@ -563,7 +585,7 @@ func TestWorkingTreeDiscardAllDirChanges(t *testing.T) {
 func TestWorkingTreeDiscardUnstagedDirChanges(t *testing.T) {
 	type scenario struct {
 		testName             string
-		node                 *testNode
+		nodes                []IFileNode
 		runner               *oscommands.FakeCmdObjRunner
 		expectedRemovedFiles []string
 	}
@@ -571,14 +593,14 @@ func TestWorkingTreeDiscardUnstagedDirChanges(t *testing.T) {
 	scenarios := []scenario{
 		{
 			testName: "directory node: removes untracked files and checks out tracked files by path, not by directory",
-			node: &testNode{
+			nodes: []IFileNode{&testNode{
 				path: "dir",
 				files: []*models.File{
 					{Path: "dir/tracked1.txt", Tracked: true},
 					{Path: "dir/tracked2.txt", Tracked: true},
 					{Path: "dir/new.txt", Tracked: false},
 				},
-			},
+			}},
 			// Must checkout the individual files, not "dir" — otherwise a filter would be ignored.
 			runner: oscommands.NewFakeRunner(t).
 				ExpectGitArgs([]string{"checkout", "--", "dir/tracked1.txt", "dir/tracked2.txt"}, "", nil),
@@ -586,32 +608,53 @@ func TestWorkingTreeDiscardUnstagedDirChanges(t *testing.T) {
 		},
 		{
 			testName: "directory node: staged-but-not-committed file (Tracked=false, HasStagedChanges=true) is left alone; purely untracked file is removed",
-			node: &testNode{
+			nodes: []IFileNode{&testNode{
 				path: "dir",
 				files: []*models.File{
-					// Staged new files: not removed from disk (RemoveUntrackedDirFiles
-					// skips staged files), but checked out in case it also has
-					// unstaged changes on top (AM status).
+					// Staged new files: not removed from disk, but checked out in
+					// case they also have unstaged changes on top (AM status).
 					{Path: "dir/staged-new1.txt", Tracked: false, Added: true, HasStagedChanges: true},
 					{Path: "dir/staged-new2.txt", Tracked: false, Added: true, HasStagedChanges: true},
 					// Purely untracked file: removed from disk, not checked out.
 					{Path: "dir/untracked.txt", Tracked: false, Added: true, HasStagedChanges: false},
 				},
-			},
+			}},
 			runner: oscommands.NewFakeRunner(t).
 				ExpectGitArgs([]string{"checkout", "--", "dir/staged-new1.txt", "dir/staged-new2.txt"}, "", nil),
 			expectedRemovedFiles: []string{"dir/untracked.txt"},
 		},
 		{
 			testName: "file node: added and unstaged file is removed from disk",
-			node: &testNode{
+			nodes: []IFileNode{&testNode{
 				path:  "new.txt",
 				files: []*models.File{{Path: "new.txt", Added: true}},
 				file:  &models.File{Path: "new.txt", Added: true, HasStagedChanges: false},
-			},
+			}},
 			runner:               oscommands.NewFakeRunner(t),
 			expectedRemovedFiles: []string{"new.txt"},
 		},
+		{
+			testName: "files from multiple nodes are batched into a single checkout call",
+			nodes: []IFileNode{
+				&testNode{
+					path: "dir1",
+					files: []*models.File{
+						{Path: "dir1/tracked.txt", Tracked: true},
+						{Path: "dir1/untracked.txt", Tracked: false},
+					},
+				},
+				&testNode{
+					path: "dir2",
+					files: []*models.File{
+						{Path: "dir2/tracked.txt", Tracked: true},
+						{Path: "dir2/untracked.txt", Tracked: false},
+					},
+				},
+			},
+			runner: oscommands.NewFakeRunner(t).
+				ExpectGitArgs([]string{"checkout", "--", "dir1/tracked.txt", "dir2/tracked.txt"}, "", nil),
+			expectedRemovedFiles: []string{"dir1/untracked.txt", "dir2/untracked.txt"},
+		},
 	}
 
 	for _, s := range scenarios {
@@ -622,7 +665,7 @@ func TestWorkingTreeDiscardUnstagedDirChanges(t *testing.T) {
 				return nil
 			}
 			instance := buildWorkingTreeCommands(commonDeps{runner: s.runner, removeFile: removeFile})
-			assert.NoError(t, instance.DiscardUnstagedDirChanges(s.node))
+			assert.NoError(t, instance.DiscardUnstagedDirChanges(s.nodes))
 			s.runner.CheckForMissingCalls()
 			assert.Equal(t, s.expectedRemovedFiles, removedFiles)
 		})

pkg/gui/controllers/files_controller.go

@@ -1496,10 +1496,9 @@ func (self *FilesController) remove(selectedNodes []*filetree.FileNode) error {
 				defer self.context().CancelRangeSelect()
 			}
 
-			for _, node := range selectedNodes {
-				if err := self.c.Git().WorkingTree.DiscardAllDirChanges(node); err != nil {
-					return err
-				}
+			nodes := lo.Map(selectedNodes, func(n *filetree.FileNode, _ int) git_commands.IFileNode { return n })
+			if err := self.c.Git().WorkingTree.DiscardAllDirChanges(nodes); err != nil {
+				return err
 			}
 
 			self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC, Scope: []types.RefreshableView{types.FILES, types.WORKTREES}})
@@ -1523,10 +1522,9 @@ func (self *FilesController) remove(selectedNodes []*filetree.FileNode) error {
 				defer self.context().CancelRangeSelect()
 			}
 
-			for _, node := range selectedNodes {
-				if err := self.c.Git().WorkingTree.DiscardUnstagedDirChanges(node); err != nil {
-					return err
-				}
+			nodes := lo.Map(selectedNodes, func(n *filetree.FileNode, _ int) git_commands.IFileNode { return n })
+			if err := self.c.Git().WorkingTree.DiscardUnstagedDirChanges(nodes); err != nil {
+				return err
 			}
 
 			self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC, Scope: []types.RefreshableView{types.FILES, types.WORKTREES}})