Stage affected unstaged files when applying or reverting a patch

Unlike moving a patch to the index, applying or reverting a patch didn't
auto-stash, which means that applying a patch when there's a modified (but
unstaged) file in the working tree would error out with the message "error:
file1: does not match index", regardless of whether those modifications conflict
with the patch or not.

To fix this, we *could* add auto-stashing like we do for the "move patch to
index" command. However, in this case we rather simply stage the affected files
(after asking for confirmation). This has a few advantages:

- it only changes the staging state of those files that are contained in the
patch (whereas auto-stashing always changes all files to unstaged)
- it doesn't unnecessarily show a confirmation if none of the modified files are
affected by the patch
- if the patch conflicts with the modified files, the conflicts were "backwards"
("ours" was the patch, "theirs" the modified file); it is more logical if "ours"
is the current state of the file, and "theirs" is the patch.

It's a little unfortunate that the behavior isn't exactly the same as for "move
patch to index", but for that one we do need the auto-stash because of the
rebase that runs behind the scenes.

Author: stefanhaller

pkg/gui/controllers/custom_patch_options_menu_action.go

@@ -4,10 +4,13 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/jesseduffield/generics/set"
 	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/controllers/helpers"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/samber/lo"
 )
 
 type CustomPatchOptionsMenuAction struct {
@@ -267,16 +270,42 @@ func (self *CustomPatchOptionsMenuAction) handlePullPatchIntoNewCommitBefore() e
 func (self *CustomPatchOptionsMenuAction) handleApplyPatch(reverse bool) error {
 	self.returnFocusFromPatchExplorerIfNecessary()
 
-	action := self.c.Tr.Actions.ApplyPatch
-	if reverse {
-		action = "Apply patch in reverse"
+	affectedUnstagedFiles := self.getAffectedUnstagedFiles()
+
+	apply := func() error {
+		action := self.c.Tr.Actions.ApplyPatch
+		if reverse {
+			action = "Apply patch in reverse"
+		}
+		self.c.LogAction(action)
+
+		if len(affectedUnstagedFiles) > 0 {
+			if err := self.c.Git().WorkingTree.StageFiles(affectedUnstagedFiles, nil); err != nil {
+				return err
+			}
+		}
+
+		if err := self.c.Git().Patch.ApplyCustomPatch(reverse, true); err != nil {
+			return err
+		}
+
+		self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC})
+		return nil
 	}
-	self.c.LogAction(action)
-	if err := self.c.Git().Patch.ApplyCustomPatch(reverse, true); err != nil {
-		return err
+
+	if len(affectedUnstagedFiles) > 0 {
+		self.c.Confirm(types.ConfirmOpts{
+			Title:  self.c.Tr.MustStageFilesAffectedByPatchTitle,
+			Prompt: self.c.Tr.MustStageFilesAffectedByPatchWarning,
+			HandleConfirm: func() error {
+				return apply()
+			},
+		})
+
+		return nil
 	}
-	self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC})
-	return nil
+
+	return apply()
 }
 
 func (self *CustomPatchOptionsMenuAction) copyPatchToClipboard() error {
@@ -291,3 +320,17 @@ func (self *CustomPatchOptionsMenuAction) copyPatchToClipboard() error {
 
 	return nil
 }
+
+// Returns a list of files that have unstaged changes and are contained in the patch.
+func (self *CustomPatchOptionsMenuAction) getAffectedUnstagedFiles() []string {
+	unstagedFiles := set.NewFromSlice(lo.FilterMap(self.c.Model().Files, func(f *models.File, _ int) (string, bool) {
+		if f.GetHasUnstagedChanges() {
+			return f.GetPath(), true
+		}
+		return "", false
+	}))
+
+	return lo.Filter(self.c.Git().Patch.PatchBuilder.AllFilesInPatch(), func(patchFile string, _ int) bool {
+		return unstagedFiles.Includes(patchFile)
+	})
+}

pkg/i18n/english.go

@@ -820,6 +820,8 @@ type TranslationSet struct {
 	MovePatchToSelectedCommit                string
 	MovePatchToSelectedCommitTooltip         string
 	CopyPatchToClipboard                     string
+	MustStageFilesAffectedByPatchTitle       string
+	MustStageFilesAffectedByPatchWarning     string
 	NoMatchesFor                             string
 	MatchesFor                               string
 	SearchKeybindings                        string
@@ -1909,6 +1911,8 @@ func EnglishTranslationSet() *TranslationSet {
 		MovePatchToSelectedCommit:                "Move patch to selected commit (%s)",
 		MovePatchToSelectedCommitTooltip:         "Move the patch out of its original commit and into the selected commit. This is achieved by starting an interactive rebase at the original commit, applying the patch in reverse, then continuing the rebase up to the selected commit, before applying the patch forward and amending the selected commit. The rebase is then continued to completion. If commits between the source and destination commit depend on the patch, you may need to resolve conflicts.",
 		CopyPatchToClipboard:                     "Copy patch to clipboard",
+		MustStageFilesAffectedByPatchTitle:       "Must stage files",
+		MustStageFilesAffectedByPatchWarning:     "Applying a patch to the index requires staging the unstaged files that are affected by the patch. Note that you might get conflicts when applying the patch. Continue?",
 		NoMatchesFor:                             "No matches for '%s' %s",
 		ExitSearchMode:                           "%s: Exit search mode",
 		ExitTextFilterMode:                       "%s: Exit filter mode",

pkg/integration/tests/patch_building/apply_with_modified_file_conflict.go

@@ -53,8 +53,31 @@ var ApplyWithModifiedFileConflict = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Common().SelectPatchOption(MatchesRegexp(`Apply patch$`))
 
+		t.ExpectPopup().Confirmation().Title(Equals("Must stage files")).
+			Content(Contains("Applying a patch to the index requires staging the unstaged files that are affected by the patch.")).
+			Confirm()
+
 		t.ExpectPopup().Alert().Title(Equals("Error")).
-			Content(Equals("error: file1: does not match index")).
+			Content(Contains("Applied patch to 'file1' with conflicts.")).
 			Confirm()
+
+		t.Views().Files().
+			Focus().
+			Lines(
+				Equals("UU file1").IsSelected(),
+			).
+			PressEnter()
+
+		t.Views().MergeConflicts().
+			IsFocused().
+			Lines(
+				Equals("<<<<<<< ours"),
+				Equals("111"),
+				Equals("======="),
+				Equals("11"),
+				Equals(">>>>>>> theirs"),
+				Equals("2"),
+				Equals("3"),
+			)
 	},
 })

pkg/integration/tests/patch_building/apply_with_modified_file_no_conflict.go

@@ -53,8 +53,17 @@ var ApplyWithModifiedFileNoConflict = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Common().SelectPatchOption(MatchesRegexp(`Apply patch$`))
 
-		t.ExpectPopup().Alert().Title(Equals("Error")).
-			Content(Equals("error: file1: does not match index")).
+		t.ExpectPopup().Confirmation().Title(Equals("Must stage files")).
+			Content(Contains("Applying a patch to the index requires staging the unstaged files that are affected by the patch.")).
 			Confirm()
+
+		t.Views().Files().
+			Focus().
+			Lines(
+				Equals("M  file1").IsSelected(),
+			)
+
+		t.Views().Main().
+			Content(Contains("-1\n+11\n 2\n 3\n+4"))
 	},
 })