Advance the selection to the next change after staging from the main view

Staging or unstaging a selection re-renders the focused main view's diff (the
staged lines move to the other side), which left the selection at a now-stale
view position — often off the new, shorter content, so it looked like there was
no selection at all.

Install a restore before the re-render that lands the selection on the change
nearest the one just acted on, the same way the staging view advances: the
change block after the selection, else the one before it, else the acted-on line
itself (which only survives when the whole side was staged and the view flips to
the other side). It re-selects in the current mode, so hunk staging walks hunk to
hunk. A range selection collapses to a single line first, since staging consumes
it. This rides the existing restore-by-identity machinery (as the escape restore
does), so it works over any conforming rendering.

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

pkg/gui/controllers/helpers/staging_helper.go

@@ -266,6 +266,40 @@ func (self *StagingHelper) RestoreFocusedMainViewOnEscape(explorerView, mainView
 	})
 }
 
+// RevealSelectionAfterStaging arranges, after staging or unstaging re-renders the
+// focused main view's diff, for the selection to land on the change nearest the one
+// just acted on rather than at a stale (and possibly off-content) position — the same
+// "advance to the next change" the staging view does. firstLine and lastLine bound
+// the just-staged selection in the current (pre-staging) diff; place positions and
+// re-selects the landed line in the current select mode. Install it before triggering
+// the re-render: it rides the next render of view, like RestoreFocusedMainViewOnEscape.
+//
+// The candidates, in priority order, are the change block after the selection, then
+// the one before it (both survive staging — only the selection itself was staged),
+// then the selection's own first line, which only turns up when the whole side was
+// staged and the view flips to show the other (staged) side. If none survives, place
+// isn't called and the view re-renders without moving the selection.
+func (self *StagingHelper) RevealSelectionAfterStaging(view *gocui.View, firstLine int, lastLine int, place func(viewLine int)) {
+	var candidates []diffLineAnchor
+	addByViewLine := func(viewLine int) {
+		if info, ok := self.GetDiffLineInfoForView(view, viewLine); ok {
+			candidates = append(candidates, diffLineAnchor{identity: info})
+		}
+	}
+
+	if next, ok := self.AdjacentChangeBlock(view, lastLine, true); ok {
+		addByViewLine(next)
+	}
+	if prev, ok := self.AdjacentChangeBlock(view, firstLine, false); ok {
+		addByViewLine(prev)
+	}
+	addByViewLine(firstLine)
+
+	self.restoreDiffLinePositionOnRerender(view, candidates, func(_ diffLineAnchor, viewLine int) {
+		place(viewLine)
+	})
+}
+
 // diffLineAnchor is a candidate line for a restore to land on after a re-render: a
 // patch identity to find the line by, plus the screen row (offset from the view's
 // top) it was on when captured, for a restore that wants to put it back there.

pkg/gui/controllers/main_view_controller.go

@@ -281,7 +281,30 @@ func (self *MainViewController) stageSelectedLine() error {
 	if handler == nil {
 		return nil
 	}
-	first, last := self.context.GetView().SelectedLineRange()
+
+	v := self.context.GetView()
+	first, last := v.SelectedLineRange()
+
+	// Staging consumes the selected range, so a range selection collapses back to a
+	// single line; hunk mode stays on, to land on the next hunk.
+	sel := self.sel()
+	if sel.Mode == context.DiffSelectModeRange {
+		sel.Mode = context.DiffSelectModeLine
+		sel.RangeIsSticky = false
+	}
+
+	// The diff re-renders after staging; install a restore (before the handler
+	// triggers that re-render) so the selection lands on the next change rather than
+	// at a now-meaningless position.
+	self.c.Helpers().Staging.RevealSelectionAfterStaging(v, first, last, func(viewLine int) {
+		if self.sel().Mode == context.DiffSelectModeHunk {
+			self.selectHunkAround(viewLine)
+		} else {
+			v.CancelRangeSelect()
+			showSelectionAtLine(v, viewLine, true)
+		}
+	})
+
 	return handler(self.context.GetViewName(), first, last)
 }
 

pkg/integration/tests/staging/select_next_hunk_after_staging_from_main_view.go

@@ -0,0 +1,48 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var SelectNextHunkAfterStagingFromMainView = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "After staging a hunk from the focused main view, the selection advances to the next hunk rather than getting lost",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = true
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "one\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten\n")
+		shell.Commit("one")
+
+		// Two separate change blocks, far enough apart to stay distinct hunks.
+		shell.UpdateFile("file1", "one\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)
+
+		t.Views().Main().
+			IsFocused().
+			SelectedLines(
+				Contains("-three"),
+				Contains("+THREE"),
+			).
+			PressPrimaryAction().
+			Tap(func() {
+				t.Views().Secondary().
+					ContainsLines(
+						Contains("-three"),
+						Contains("+THREE"),
+					)
+			}).
+			// The selection didn't get lost: it advanced to the next (and now only
+			// remaining) hunk, which is what the unstaged half shows.
+			SelectedLines(
+				Contains("-nine"),
+				Contains("+NINE"),
+			)
+	},
+})

pkg/integration/tests/staging/select_next_hunk_after_unstaging_from_main_view.go

@@ -0,0 +1,52 @@
+package staging
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var SelectNextHunkAfterUnstagingFromMainView = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "After unstaging a hunk from the staged half of the focused main view, the selection advances to the next staged hunk",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Gui.UseHunkModeInStagingView = true
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "one\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten\neleven\n")
+		shell.Commit("one")
+
+		// Two staged hunks...
+		shell.UpdateFileAndAdd("file1", "one\ntwo\nTHREE\nfour\nfive\nsix\nseven\neight\nNINE\nten\neleven\n")
+		// ...plus an unstaged change, so the main view splits into staged/unstaged.
+		shell.UpdateFile("file1", "one\ntwo\nTHREE\nfour\nfive\nSIX\nseven\neight\nNINE\nten\neleven\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Press(keys.Universal.FocusMainView)
+
+		// The unstaged half is focused first; switch to the staged half.
+		t.Views().Main().
+			IsFocused().
+			SelectedLines(
+				Contains("-six"),
+				Contains("+SIX"),
+			).
+			Press(keys.Universal.TogglePanel)
+
+		t.Views().Secondary().
+			IsFocused().
+			SelectedLines(
+				Contains("-three"),
+				Contains("+THREE"),
+			).
+			// Unstage the first staged hunk.
+			PressPrimaryAction().
+			// The selection advances to the next staged hunk rather than getting lost.
+			SelectedLines(
+				Contains("-nine"),
+				Contains("+NINE"),
+			)
+	},
+})

pkg/integration/tests/test_list.go

@@ -411,6 +411,8 @@ var tests = []*components.IntegrationTest{
 	staging.DiscardAllChanges,
 	staging.Search,
 	staging.SelectHunkOnFocusingMainView,
+	staging.SelectNextHunkAfterStagingFromMainView,
+	staging.SelectNextHunkAfterUnstagingFromMainView,
 	staging.SelectNextLineAfterStagingInTwoHunkDiff,
 	staging.SelectNextLineAfterStagingIsolatedAddedLine,
 	staging.StageHunkFromMainView,