Carry the selection to the next change after staging

Staging takes the lines it acted on out of the diff, so the selection has
nothing to sit on afterwards and would be left wherever those lines used to
be. What the user wants is the change that moved up into their place, so
that pressing the key again goes on to the next one — which is how staging
line by line through a file works.

The line acted on is gone, so it can't be remembered by identity the way a
re-render of the same diff remembers one; what is remembered instead is its
place in the sequence of the diff's changes, which the change after it
inherits. Staging the last change is the one case with nothing to inherit
it, and there the selection stays on the last change there is.

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

pkg/gui/controllers/helpers/diff_line_helper.go

@@ -86,6 +86,18 @@ func (self *DiffLineHelper) diffLineIdentitiesFromRecords(metadata []string) []t
 	return self.diffLineInfos(parseDiffLineRecords(metadata))
 }
 
+// diffLineInfoFromRecords recovers a row's own identity from the records the diff
+// renderer stated for it. That is the line the row leads with, of the ones
+// diffLineIdentitiesFromRecords finds on it. ok is false when the row carries no
+// record we understand.
+func (self *DiffLineHelper) diffLineInfoFromRecords(metadata []string) (types.DiffLineInfo, bool) {
+	identities := self.diffLineIdentitiesFromRecords(metadata)
+	if len(identities) == 0 {
+		return types.DiffLineInfo{}, false
+	}
+	return identities[0], true
+}
+
 // resolvedDiffLine is one rendered row's recovered identity, plus whether it could
 // be recovered at all — the element of the table resolveDiffLines produces.
 type resolvedDiffLine struct {

pkg/gui/controllers/helpers/diff_line_restore.go

@@ -264,6 +264,65 @@ func (self *DiffLineHelper) restoreDiffLinePositionOnRerender(
 	)
 }
 
+// ChangeLineOrdinal returns how many change lines of view's rendered diff come before
+// the given buffer line — that line's place in the sequence of changes.
+//
+// It is how a place in a diff is remembered across acting on it: an action consumes
+// the lines it acted on, so the identity of the line the user was on is gone, but the
+// place it left behind is the same one that identity used to have.
+func (self *DiffLineHelper) ChangeLineOrdinal(view *gocui.View, bufferLine int) int {
+	ordinal := 0
+	for i, row := range self.resolveDiffLines(view.DiffLineContents()) {
+		if i >= bufferLine {
+			break
+		}
+		if row.ok && row.info.IsChange() {
+			ordinal++
+		}
+	}
+	return ordinal
+}
+
+// RevealChangeLineAtOrdinal arranges for view's next re-render to be shown with the
+// change line at the given ordinal placed by place — the diff having changed under the
+// user, this is where what they were doing carries on. When the new diff has fewer
+// changes than that, because the ones acted on were its last, it lands on the last
+// change left.
+func (self *DiffLineHelper) RevealChangeLineAtOrdinal(view *gocui.View, ordinal int, place func(viewLine int)) {
+	// How many change lines the incremental search has passed, so that it can carry on
+	// counting where it left off.
+	seen := 0
+
+	self.installDiffLineRestore(view,
+		func(rows []gocui.DiffLineContent, offset int) (int, bool) {
+			for i, row := range rows {
+				if info, ok := self.diffLineInfoFromRecords(row.Metadata); ok && info.IsChange() {
+					if seen == ordinal {
+						return offset + i, true
+					}
+					seen++
+				}
+			}
+			return 0, false
+		},
+		func(contents []gocui.DiffLineContent) (int, bool) {
+			last, count := -1, 0
+			for i, row := range self.resolveDiffLines(contents) {
+				if !row.ok || !row.info.IsChange() {
+					continue
+				}
+				if count == ordinal {
+					return i, true
+				}
+				count++
+				last = i
+			}
+			return last, last != -1
+		},
+		place,
+	)
+}
+
 // installDiffLineRestore is what the restores are built on: it arranges for view's
 // next re-render to be revealed with the row a search finds in it placed by place,
 // instead of from the top.

pkg/gui/controllers/main_view_controller.go

@@ -300,6 +300,36 @@ func (self *MainViewController) primaryAction() error {
 	return actions.PrimaryAction(self.context, first, last)
 }
 
+// revealSelectionAfterAction moves the focused main view's selection to the change
+// that takes the place of the one just acted on, once the changed diff has re-rendered.
+// Call it from the panel's action handler with the pane it acted in and the buffer line
+// the selection starts on, before triggering the re-render.
+//
+// The line acted on is gone from the diff, so what is remembered is its place among the
+// diff's changes: the next change moves up into it, which is where you want to be to
+// carry on. A range collapses to a single line at its start, and hunk mode selects the
+// whole block it lands in, so that pressing the key again acts on the next hunk.
+func revealSelectionAfterAction(c *ControllerCommon, pane types.DiffPaneContext, firstBufferLine int) {
+	view := pane.GetView()
+	ordinal := c.Helpers().DiffLine.ChangeLineOrdinal(view, firstBufferLine)
+
+	sel := pane.DiffSelectState()
+	if sel.Mode == types.DiffSelectModeRange {
+		sel.Mode = types.DiffSelectModeLine
+		sel.RangeIsSticky = false
+	}
+	selectHunk := sel.Mode == types.DiffSelectModeHunk
+
+	c.Helpers().DiffLine.RevealChangeLineAtOrdinal(view, ordinal, func(viewLine int) {
+		if selectHunk {
+			selectDiffHunk(c, pane, viewLine, true)
+			return
+		}
+		view.CancelRangeSelect()
+		showSelectionAtLine(view, viewLine, true)
+	})
+}
+
 // workingTreeActionDescription gives a command's description only where the command
 // applies — over the working tree's diff — so that it is listed there and nowhere else.
 func (self *MainViewController) workingTreeActionDescription(description string) func() string {
@@ -622,12 +652,12 @@ func (self *MainViewController) selectHunkAround(changeViewLine int, scrollIntoV
 // to go there; a click leaves it false, so that the view doesn't move under the mouse
 // when the block the click landed in starts above the viewport.
 func selectDiffHunk(
-	c *ControllerCommon, mainContext *context.MainContext, changeViewLine int, scrollIntoView bool,
+	c *ControllerCommon, pane types.DiffPaneContext, changeViewLine int, scrollIntoView bool,
 ) {
-	view := mainContext.GetView()
+	view := pane.GetView()
 	start, end, ok := c.Helpers().DiffLine.ChangeBlockBounds(view, changeViewLine)
 	if !ok {
-		mainContext.DiffSelectState().Mode = types.DiffSelectModeLine
+		pane.DiffSelectState().Mode = types.DiffSelectModeLine
 		view.CancelRangeSelect()
 		showSelectionAtLine(view, changeViewLine, scrollIntoView)
 		return

pkg/gui/controllers/working_tree_diff_actions.go

@@ -66,7 +66,7 @@ func (self *WorkingTreeDiffActions) PrimaryAction(pane types.DiffPaneContext, fi
 
 	// Either way the patch goes to the index: forwards from the unstaged side to stage
 	// it, backwards from the staged side to take it back out.
-	return self.applyDiffLineSelection(infos, onStagedSide,
+	return self.applyDiffLineSelection(pane, firstBufferLine, infos, onStagedSide,
 		git_commands.ApplyPatchOpts{Reverse: onStagedSide, Cached: true})
 }
 
@@ -88,7 +88,10 @@ func (self *WorkingTreeDiffActions) diffLineSelection(
 // applyDiffLineSelection applies the selected change lines, a patch per file, and
 // re-renders what that changed. onStagedSide says which of the file's two diffs the
 // lines were selected in and so are to be found in; opts says how to apply them.
+// firstBufferLine is where the selection started, which is where the work carries on
+// from once the diff has changed under it.
 func (self *WorkingTreeDiffActions) applyDiffLineSelection(
+	pane types.DiffPaneContext, firstBufferLine int,
 	infos []types.DiffLineInfo, onStagedSide bool, opts git_commands.ApplyPatchOpts,
 ) error {
 	self.c.LogAction(self.c.Tr.Actions.ApplyPatch)
@@ -106,6 +109,11 @@ func (self *WorkingTreeDiffActions) applyDiffLineSelection(
 		}
 	}
 
+	// The refresh below queues the re-render of the diff we just changed; this rides it,
+	// so that the selection ends up on the change that took the place of the one acted
+	// on rather than at a position that means nothing any more.
+	revealSelectionAfterAction(self.c, pane, firstBufferLine)
+
 	// Block input until the refresh has landed, so that a quick second keypress acts on
 	// the diff as it now is rather than on the one we just changed.
 	self.c.RefreshBlockingInput(types.RefreshOptions{Scope: []types.RefreshableView{types.FILES}})

pkg/integration/tests/main_view/select_next_change_after_staging.go

@@ -0,0 +1,54 @@
+package main_view
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var SelectNextChangeAfterStaging = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "After staging from the focused main view the selection lands on the change that took the place of the one staged",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(cfg *config.AppConfig) {
+		cfg.GetUserConfig().Gui.UseHunkModeInStagingView = false
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "one\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten\n")
+		shell.Commit("one")
+
+		shell.UpdateFile("file1", "one\nADD1\nADD2\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nNINE\nten\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Press(keys.Universal.FocusMainView)
+
+		// Line by line, each press leaves the selection on the next change to stage.
+		t.Views().Main().
+			IsFocused().
+			SelectedLines(
+				Contains("+ADD1"),
+			).
+			PressPrimaryAction().
+			SelectedLines(
+				Contains("+ADD2"),
+			).
+			PressPrimaryAction().
+			SelectedLines(
+				Contains("-nine"),
+			).
+			// A hunk goes the same way: the block after the one staged takes its place,
+			// and here there is none, so the last change stays selected.
+			Press(keys.Main.ToggleSelectHunk).
+			SelectedLines(
+				Contains("-nine"),
+				Contains("+NINE"),
+			).
+			PressPrimaryAction().
+			Tap(func() {
+				t.Views().Files().Lines(
+					Contains("M  file1"),
+				)
+			})
+	},
+})

pkg/integration/tests/test_list.go

@@ -405,6 +405,7 @@ var tests = []*components.IntegrationTest{
 	main_view.SelectHunkOnFocusingMainView,
 	main_view.SelectInADiffReadInPart,
 	main_view.SelectLineWhenWholeFileIsOneHunk,
+	main_view.SelectNextChangeAfterStaging,
 	main_view.SelectVisibleChangeOnFocusingMainView,
 	main_view.SelectVisibleHunkOnFocusingMainView,
 	main_view.SelectionCommandsOnlyWhereTheyApply,