Improve performance of discarding changes in large directories

Previously it would iterate over all changed files and call git checkout or git
reset for each one, which can take forever if there are hundreds or thousands of
files. Now it batches these into a single command if possible (taking care of
still passing the individual path names to the git call rather than just the
directory, which is necessary for making it work correctly when filtering --
this was actually broken for the "Discard unstaged changes" command, which is
fixed here).

pkg/commands/git_commands/working_tree.go

@@ -185,8 +185,64 @@ type IFileNode interface {
 }
 
 func (self *WorkingTreeCommands) DiscardAllDirChanges(node IFileNode) error {
-	// this could be more efficient but we would need to handle all the edge cases
-	return node.ForEachFile(self.DiscardAllFileChanges)
+	// 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
+		}
+
+		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)
+			return nil
+		}
+
+		filesToCheckout = append(filesToCheckout, file.Path)
+		return nil
+	})
+
+	for _, file := range specialFiles {
+		if err := self.DiscardAllFileChanges(file); err != nil {
+			return err
+		}
+	}
+
+	if err := runGitCmdOnPaths("reset", filesToReset, self.cmd); err != nil {
+		return err
+	}
+
+	for _, path := range filesToRemove {
+		if err := self.os.RemoveFile(path); err != nil {
+			return err
+		}
+	}
+
+	return runGitCmdOnPaths("checkout", filesToCheckout, self.cmd)
 }
 
 func (self *WorkingTreeCommands) DiscardUnstagedDirChanges(node IFileNode) error {
@@ -196,8 +252,14 @@ func (self *WorkingTreeCommands) DiscardUnstagedDirChanges(node IFileNode) error
 			return err
 		}
 
-		cmdArgs := NewGitCmd("checkout").Arg("--", node.GetPath()).ToArgv()
-		if err := self.cmd.New(cmdArgs).Run(); err != nil {
+		// 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()
+		})
+		if err := runGitCmdOnPaths("checkout", trackedPaths, self.cmd); err != nil {
 			return err
 		}
 	} else {
@@ -215,7 +277,7 @@ func (self *WorkingTreeCommands) DiscardUnstagedDirChanges(node IFileNode) error
 
 func (self *WorkingTreeCommands) RemoveUntrackedDirFiles(node IFileNode) error {
 	untrackedFilePaths := node.GetFilePathsMatching(
-		func(file *models.File) bool { return !file.GetIsTracked() },
+		func(file *models.File) bool { return !file.GetIsTracked() && !file.GetHasStagedChanges() },
 	)
 
 	for _, path := range untrackedFilePaths {

pkg/commands/git_commands/working_tree_test.go

@@ -508,7 +508,7 @@ func TestWorkingTreeDiscardAllDirChanges(t *testing.T) {
 
 	scenarios := []scenario{
 		{
-			testName: "multiple tracked files make individual checkout calls",
+			testName: "multiple regular tracked files batched into a single checkout call",
 			node: &testNode{
 				files: []*models.File{
 					{Path: "a.txt", Tracked: true},
@@ -517,12 +517,10 @@ func TestWorkingTreeDiscardAllDirChanges(t *testing.T) {
 				},
 			},
 			runner: oscommands.NewFakeRunner(t).
-				ExpectGitArgs([]string{"checkout", "--", "a.txt"}, "", nil).
-				ExpectGitArgs([]string{"checkout", "--", "b.txt"}, "", nil).
-				ExpectGitArgs([]string{"checkout", "--", "c.txt"}, "", nil),
+				ExpectGitArgs([]string{"checkout", "--", "a.txt", "b.txt", "c.txt"}, "", nil),
 		},
 		{
-			testName: "staged files each make an individual reset then checkout",
+			testName: "staged files batched into a single reset then a single checkout",
 			node: &testNode{
 				files: []*models.File{
 					{Path: "a.txt", Tracked: true, HasStagedChanges: true},
@@ -530,10 +528,8 @@ func TestWorkingTreeDiscardAllDirChanges(t *testing.T) {
 				},
 			},
 			runner: oscommands.NewFakeRunner(t).
-				ExpectGitArgs([]string{"reset", "--", "a.txt"}, "", nil).
-				ExpectGitArgs([]string{"checkout", "--", "a.txt"}, "", nil).
-				ExpectGitArgs([]string{"reset", "--", "b.txt"}, "", nil).
-				ExpectGitArgs([]string{"checkout", "--", "b.txt"}, "", nil),
+				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",
@@ -574,7 +570,7 @@ func TestWorkingTreeDiscardUnstagedDirChanges(t *testing.T) {
 
 	scenarios := []scenario{
 		{
-			testName: "directory node: uses directory path for checkout",
+			testName: "directory node: removes untracked files and checks out tracked files by path, not by directory",
 			node: &testNode{
 				path: "dir",
 				files: []*models.File{
@@ -583,27 +579,28 @@ func TestWorkingTreeDiscardUnstagedDirChanges(t *testing.T) {
 					{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"}, "", nil),
+				ExpectGitArgs([]string{"checkout", "--", "dir/tracked1.txt", "dir/tracked2.txt"}, "", nil),
 			expectedRemovedFiles: []string{"dir/new.txt"},
 		},
 		{
-			testName: "directory node: staged-but-not-committed file (Tracked=false, HasStagedChanges=true) is removed along with untracked files",
+			testName: "directory node: staged-but-not-committed file (Tracked=false, HasStagedChanges=true) is left alone; purely untracked file is removed",
 			node: &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).
 					{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"}, "", nil),
-			// All files are removed because the predicate of GetFilePathsMatching in
-			// RemoveUntrackedDirFiles is just !Tracked. git checkout -- dir then restores the
-			// staged file from the index. This is a bit wasteful, and we'll improve it at the end
-			// of this branch.
-			expectedRemovedFiles: []string{"dir/staged-new1.txt", "dir/staged-new2.txt", "dir/untracked.txt"},
+				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",

pkg/integration/tests/file/discard_unstaged_dir_changes_when_filtering.go

@@ -56,23 +56,14 @@ var DiscardUnstagedDirChangesWhenFiltering = NewIntegrationTest(NewIntegrationTe
 			}).
 			Press(keys.Universal.Return). // Cancel filtering
 			Lines(
-				/* EXPECTED:
 				Equals("▼ dir").IsSelected(),
 				Equals("  M  file-one"),
 				Equals("  MM file-two"),
 				Equals("  ?? unstaged-file-two"),
-				ACTUAL: */
-				Equals("▼ dir").IsSelected(),
-				Equals("  M  file-one"),
-				Equals("  M  file-two"),
-				Equals("  ?? unstaged-file-two"),
 			)
 
 		t.FileSystem().FileContent("dir/file-one", Equals("original content\nnew content\n"))
-		/* EXPECTED:
 		t.FileSystem().FileContent("dir/file-two", Equals("original content\nnew content\neven newer content\n"))
-		ACTUAL: */
-		t.FileSystem().FileContent("dir/file-two", Equals("original content\nnew content\n"))
 		t.FileSystem().PathNotPresent("dir/unstaged-file-one")
 		t.FileSystem().FileContent("dir/unstaged-file-two", Equals("unstaged file"))
 	},