Keep the diff selection when a commit is rewritten under it

Removing lines from a commit, moving a custom patch out of one, undoing
either — none of it goes through the focused main view, so a selection over
that commit's diff was left where it was, painted over a rendering of a diff
that no longer exists.

Every render now asks whether it is showing a different diff than the one on
screen, which is what a rewritten commit looks like, and puts the selection
back on the change that has taken its place — the same reveal an action in
the view does for itself. It stands down for a render of the same diff,
where a selection, perhaps a range still being made, is exactly right as it
is, and for one that something more precise is already waiting to place.

The key a render is remembered under moves ahead of anything that might
change the command's arguments, so that it says which diff is being
rendered and nothing else.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

pkg/gui/controllers/helpers/diff_line_restore.go

@@ -563,15 +563,17 @@ func patchLineOf(info types.DiffLineInfo) patchLine {
 // in another pane — and shows no selection until the restore places one, so that what
 // it was left showing the last time it was used doesn't appear for a frame.
 //
-// done is called once the selection is where it belongs, or once it turns out that no
-// render is coming to put it there, for a caller that must not let the user act again
-// in between.
+// done, which may be nil, is called once the selection is where it belongs, or once it
+// turns out that no render is coming to put it there — for a caller that must not let
+// the user act again in between.
 func (self *DiffLineHelper) RevealSelectionAfterAction(
 	source types.DiffPaneContext, target types.DiffPaneContext, firstLineIdx int, done func(),
 ) {
 	ordinal, ok := self.ChangeLineOrdinal(source.GetView(), firstLineIdx)
 	if !ok {
-		done()
+		if done != nil {
+			done()
+		}
 		return
 	}
 

pkg/gui/main_panels.go

@@ -1,6 +1,8 @@
 package gui
 
 import (
+	"strings"
+
 	"github.com/jesseduffield/lazygit/pkg/gocui"
 	"github.com/jesseduffield/lazygit/pkg/gui/context"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
@@ -118,6 +120,7 @@ func (gui *Gui) refreshMainViews(opts types.RefreshMainOpts) {
 	// Before the render is triggered, so that the pane the focus moves into can be
 	// told where to put its selection as it renders.
 	gui.followFocusIntoWorkablePane(opts)
+	gui.keepDiffSelectionAcrossACommitRewrite(opts)
 
 	gui.moveMainContextPairToTop(opts.Pair)
 
@@ -244,6 +247,71 @@ func (gui *Gui) followFocusIntoWorkablePane(opts types.RefreshMainOpts) {
 	gui.State.ContextMgr.Push(target, types.OnFocusOpts{})
 }
 
+// keepDiffSelectionAcrossACommitRewrite arranges for a selection in the focused main
+// view to come back on the same change of the diff when the render about to happen is
+// of a different diff from the one on screen. This happens when a commit is rewritten
+// under the user, by moving a patch out of it, discarding lines from it, or undoing
+// either. The selection is then left at a position in a rendering that no longer exists.
+//
+// What is remembered is which change of the diff the selection was on rather than which
+// line of which file, a rewrite being precisely a change to those lines: the change that
+// takes its place is where the work carries on.
+//
+// It is asked of every render, and does nothing unless all three of these hold: there
+// is a selection to keep; nothing more precise is already waiting to be put back (the
+// position preserves and the post-action reveals know better where their selection
+// belongs); and the diff really is another one. A plain refresh re-renders the same
+// diff, where the selection, possibly a range the user is in the middle of making, is
+// still exactly right.
+func (gui *Gui) keepDiffSelectionAcrossACommitRewrite(opts types.RefreshMainOpts) {
+	// The focused main view's two panes only: no other pair has a diff selection.
+	if opts.Pair.Main.GetKey() != context.NORMAL_MAIN_CONTEXT_KEY {
+		return
+	}
+
+	current := gui.State.ContextMgr.CurrentStatic().GetKey()
+	for _, pane := range []struct {
+		context types.Context
+		update  *types.ViewUpdateOpts
+	}{
+		{opts.Pair.Main, opts.Main},
+		{opts.Pair.Secondary, opts.Secondary},
+	} {
+		if pane.update == nil || pane.context.GetKey() != current {
+			continue
+		}
+		mainContext := gui.mainContextForView(pane.context.GetView())
+		if mainContext == nil || !mainContext.GetView().Highlight {
+			continue
+		}
+		manager := gui.getViewBufferManagerForView(mainContext.GetView())
+		if manager == nil || manager.HasRestoreForNextTask() {
+			continue
+		}
+		key, ok := diffTaskCommandKey(pane.update.Task)
+		if !ok || key == manager.GetTaskKey() {
+			continue
+		}
+
+		first, _ := mainContext.GetView().SelectedLineRange()
+		gui.helpers.DiffLine.RevealSelectionAfterAction(mainContext, mainContext, first, nil)
+	}
+}
+
+// diffTaskCommandKey returns the key the given render will be remembered under. Two
+// renders of the same diff have the same key, so comparing keys says whether a render
+// is of the diff already on screen. ok is false for a render that is a message rather
+// than a diff.
+func diffTaskCommandKey(task types.UpdateTask) (string, bool) {
+	switch task := task.(type) {
+	case *types.RunCommandTask:
+		return strings.Join(task.Cmd.Args, " "), true
+	case *types.RunDiffRendererTask:
+		return strings.Join(task.Cmd.Args, " "), true
+	}
+	return "", false
+}
+
 // onlyWorkablePane returns the main pane a render leaves as the only one worth having
 // the focus in, or nil when that is true of both of them or of neither. Being shown is
 // not the same as being worth working in: a pane the layout keeps around for the sake

pkg/gui/main_view_render.go

@@ -51,6 +51,11 @@ func (gui *Gui) newRenderTask(view *gocui.View, cmd *exec.Cmd, prefix string) er
 		return gui.newCmdTask(view, cmd, prefix)
 	}
 
+	// The key the render is remembered under says which diff it is of, so that a
+	// re-render of the same diff can be told from a render of another one. Take
+	// it before anything else can touch the command's arguments.
+	cmdStr := strings.Join(cmd.Args, " ")
+
 	// Mark the view as loading synchronously now, before the layout pass: the
 	// actual task is created in afterLayout (below), which runs after layout, so
 	// without this the next layout pass would clamp the scroll position to the
@@ -81,8 +86,6 @@ func (gui *Gui) newRenderTask(view *gocui.View, cmd *exec.Cmd, prefix string) er
 			gui.c.ErrorToast(err.Error())
 		}
 
-		cmdStr := strings.Join(cmd.Args, " ")
-
 		// This communicates to diff renderers that we're in a very simple
 		// terminal that they should not expect to have much capabilities.
 		// Moving the cursor, clearing the screen, or querying for colors are among such "advanced" capabilities.

pkg/integration/tests/patch_building/keep_selection_after_moving_patch_out_main_view.go

@@ -0,0 +1,84 @@
+package patch_building
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var KeepSelectionAfterMovingPatchOutMainView = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Moving a custom patch out of a commit leaves the focused main view's selection on a change that is still there, rather than painted over the diff the rewrite left behind",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = false
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "one\ntwo\nthree\nfour\nfive\n")
+		shell.Commit("first commit")
+
+		shell.UpdateFileAndAdd("file1", "ONE\ntwo\nTHREE\nfour\nFIVE\n")
+		shell.Commit("commit to move a patch out of")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Commits().
+			Focus().
+			Lines(
+				Contains("commit to move a patch out of").IsSelected(),
+				Contains("first commit"),
+			).
+			PressEnter()
+
+		// Take the first of the commit's three changed lines into a custom patch.
+		t.Views().CommitFiles().
+			IsFocused().
+			PressEnter()
+
+		t.Views().PatchBuilding().
+			IsFocused().
+			SelectedLines(
+				Contains("-one"),
+			).
+			Press(keys.Universal.ToggleRangeSelect).
+			Press(keys.Universal.NextItem).
+			SelectedLines(
+				Contains("-one"),
+				Contains("+ONE"),
+			).
+			PressPrimaryAction().
+			Press(keys.Universal.Return)
+
+		// Leave a range selected over the diff, spanning the lines the patch holds. The
+		// patch move doesn't go through the main view at all, so without a net nothing
+		// would move this selection off lines that the rewrite takes away.
+		t.Views().CommitFiles().
+			IsFocused().
+			Press(keys.Universal.FocusMainView)
+
+		t.Views().Main().
+			IsFocused().
+			SelectedLines(
+				Contains("-one"),
+			).
+			Press(keys.Universal.ToggleRangeSelect).
+			NavigateToLine(Contains("+THREE")).
+			SelectedLines(
+				Contains("-one"),
+				Contains("+ONE"),
+				Contains(" two"),
+				Contains("-three"),
+				Contains("+THREE"),
+			)
+
+		t.Common().SelectPatchOption(Contains("Move patch out into index"))
+
+		// The moved lines are gone from the commit, so the range collapses onto the
+		// change that has taken its place — the same place in the diff's changes, which
+		// is where the user was.
+		t.Views().Main().
+			IsFocused().
+			Content(DoesNotContain("+ONE")).
+			SelectedLines(
+				Contains("-three"),
+			)
+	},
+})

pkg/integration/tests/test_list.go

@@ -453,6 +453,7 @@ var tests = []*components.IntegrationTest{
 	patch_building.CopyRenamedFileDiff,
 	patch_building.DiscardLinesFromCommit,
 	patch_building.EditLineInPatchBuildingPanel,
+	patch_building.KeepSelectionAfterMovingPatchOutMainView,
 	patch_building.MoveRangeToIndex,
 	patch_building.MoveToEarlierCommit,
 	patch_building.MoveToEarlierCommitFromAddedFile,