Carry the acted-on lines into the pane the work lands in

Each side of a file's diff has a pane of its own, but a pane is only shown
while its side has something in it: staging the last unstaged change takes
the upper pane away, and unstaging the last staged one takes the lower one
away. The focus follows into the pane that is left, and the selection has to
be waiting there when it arrives — on the lines just acted on, which is where
they are now, unless they were discarded rather than moved, in which case on
what is left of the file. The pane being moved to shows no selection until
the restore places one, so that the selection it was left with the last time
it was used doesn't appear for a frame.

Whether the acted-on side still holds anything is a question the model can't
answer: a refresh only queues its update, and by the time it lands the
re-render this has to ride is already under way. So the answer is worked out
from what we just did — the files we changed report whether the selection
covered all of their changes, and the ones it didn't touch are as the model
describes them.

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

pkg/gui/controllers/helpers/diff_line_restore.go

@@ -346,7 +346,10 @@ func (self *DiffLineHelper) installDiffLineRestore(
 	findComplete func(contents []gocui.DiffLineContent) (int, bool),
 	place func(viewLine int),
 ) {
-	manager := self.c.GetViewBufferManagerForView(view)
+	// Get-or-create, because the pane may not have rendered anything yet: a file whose
+	// diff has only just become split has a second pane whose first render is the one
+	// this restore is for.
+	manager := self.c.GetOrCreateViewBufferManagerForView(view)
 	if manager == nil {
 		return
 	}

pkg/gui/controllers/main_view_controller.go

@@ -297,36 +297,47 @@ 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 first line
-// of the selection, before triggering the re-render.
+// revealSelectionAfterAction moves the 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, the pane the work carries on in, and the
+// first line of the selection, 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, firstLineIdx int) {
-	view := pane.GetView()
-	ordinal, ok := c.Helpers().DiffLine.ChangeLineOrdinal(view, firstLineIdx)
+// whole block it lands in, so that pressing the key again acts on the next hunk. The
+// target pane inherits that select mode, this being the same piece of work continuing
+// 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.
+func revealSelectionAfterAction(
+	c *ControllerCommon, source types.DiffPaneContext, target types.DiffPaneContext, firstLineIdx int,
+) {
+	ordinal, ok := c.Helpers().DiffLine.ChangeLineOrdinal(source.GetView(), firstLineIdx)
 	if !ok {
 		return
 	}
 
-	sel := pane.DiffSelectState()
+	sel := source.DiffSelectState()
 	if sel.Mode == types.DiffSelectModeRange {
 		sel.Mode = types.DiffSelectModeLine
 		sel.RangeIsSticky = false
 	}
+	*target.DiffSelectState() = *sel
 	selectHunk := sel.Mode == types.DiffSelectModeHunk
 
-	c.Helpers().DiffLine.RevealChangeLineAtOrdinal(view, ordinal, func(viewLine int) {
+	targetView := target.GetView()
+	if target != source {
+		target.SetHasSelectableContent(false)
+		c.Context().UpdateSelectionHighlights()
+	}
+
+	c.Helpers().DiffLine.RevealChangeLineAtOrdinal(targetView, ordinal, func(viewLine int) {
 		if selectHunk {
-			c.Helpers().DiffLine.SelectChangeBlock(pane, viewLine, true)
+			c.Helpers().DiffLine.SelectChangeBlock(target, viewLine, true)
 			return
 		}
-		view.CancelRangeSelect()
-		c.Helpers().DiffLine.ShowSelectionAtLine(view, viewLine, true)
+		targetView.CancelRangeSelect()
+		c.Helpers().DiffLine.ShowSelectionAtLine(targetView, viewLine, true)
 	})
 }
 

pkg/gui/controllers/working_tree_diff_actions.go

@@ -99,20 +99,35 @@ func (self *WorkingTreeDiffActions) applyDiffLineSelection(
 	// A directory's diff spans several files, and a patch is of one file, so the
 	// selected lines are grouped by the file they belong to and applied file by file.
 	infosByFile := lo.GroupBy(infos, func(info types.DiffLineInfo) string { return info.Path })
+	acted := set.New[string]()
+	actedSideRemains := false
 	for path, fileInfos := range infosByFile {
 		file := self.fileForDiffLinePath(path)
 		if file == nil {
 			continue
 		}
-		if err := self.applyDiffLines(file, fileInfos, onStagedSide, opts); err != nil {
+		changesLeft, err := self.applyDiffLines(file, fileInfos, onStagedSide, opts)
+		if err != nil {
 			return err
 		}
+		acted.Add(file.GetPath())
+		actedSideRemains = actedSideRemains || changesLeft
 	}
+	if !actedSideRemains {
+		actedSideRemains = self.anyFileHasChangesOnSide(acted, onStagedSide)
+	}
+
+	// Whether the other side has anything decides which pane the work carries on in.
+	// If the lines were staged, they are in the index now, so that side has them. If
+	// they were discarded, the other side was not touched, so the model still
+	// describes it correctly.
+	otherSideHasChanges := opts.Cached || self.anyFileHasChangesOnSide(set.New[string](), !onStagedSide)
 
 	// 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, firstLineIdx)
+	// on rather than at a position that means nothing any more — and in the pane the
+	// work carries on in, which is not always the one it was in.
+	self.revealSelectionInPaneItLandsIn(pane, firstLineIdx, actedSideRemains, otherSideHasChanges)
 
 	// 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.
@@ -145,9 +160,13 @@ func (self *WorkingTreeDiffActions) fileForDiffLinePath(path string) *models.Fil
 // position in the new file and differ only in being a deletion. Context lines are not
 // selected: a patch of the lines you picked keeps whatever context it needs around
 // them by itself.
+//
+// It reports whether the diff it read holds changes the selection didn't cover. The
+// caller uses this to tell whether the side acted on still has anything of this file
+// in it once we are done.
 func (self *WorkingTreeDiffActions) applyDiffLines(
 	file *models.File, infos []types.DiffLineInfo, sourceCached bool, opts git_commands.ApplyPatchOpts,
-) error {
+) (bool, error) {
 	parsedPatch := patch.Parse(self.c.Git().WorkingTree.WorktreeFileDiff(file, true, sourceCached))
 
 	type changeLine struct {
@@ -179,17 +198,19 @@ func (self *WorkingTreeDiffActions) applyDiffLines(
 		}
 	}
 
+	changesLeft := len(patchLineIndices) < changeLineCount(parsedPatch)
+
 	// Acting on every change of a file is acting on the file itself, and saying so is
 	// not the same as applying its diff. The diff of a deleted file is its content
 	// going away, and putting that into the index line by line leaves an empty file
 	// there rather than the deletion; the diff of an added one is its whole content,
 	// and taking that back out leaves an empty file in the index rather than an
 	// untracked one.
-	if opts.Cached && len(patchLineIndices) == changeLineCount(parsedPatch) {
+	if !changesLeft && opts.Cached {
 		if opts.Reverse {
-			return self.c.Git().WorkingTree.UnStageFile(file.Names(), file.Tracked)
+			return false, self.c.Git().WorkingTree.UnStageFile(file.Names(), file.Tracked)
 		}
-		return self.c.Git().WorkingTree.StageFile(file.GetPath())
+		return false, self.c.Git().WorkingTree.StageFile(file.GetPath())
 	}
 
 	patchToApply := parsedPatch.
@@ -200,10 +221,10 @@ func (self *WorkingTreeDiffActions) applyDiffLines(
 		}).
 		FormatPlain()
 	if patchToApply == "" {
-		return nil
+		return changesLeft, nil
 	}
 
-	return self.c.Git().Patch.ApplyPatch(patchToApply, opts)
+	return changesLeft, self.c.Git().Patch.ApplyPatch(patchToApply, opts)
 }
 
 // changeLineCount returns how many of a patch's lines are changes rather than context
@@ -213,3 +234,55 @@ func changeLineCount(p *patch.Patch) int {
 		return line.IsAddition() || line.IsDeletion()
 	})
 }
+
+// revealSelectionInPaneItLandsIn arranges for the selection to carry on where the work
+// does, which is not always the pane it was in.
+//
+// Each side of the diff has a pane of its own, so acting on one usually leaves
+// everything where it is. But a pane is only shown while its side has something in it:
+// staging the last unstaged change takes the upper pane away, and unstaging the last
+// staged one takes the lower one away. The refresh moves the focus into whichever pane
+// is left, and this puts the selection there to meet it — on the lines just acted on,
+// which are in that pane now, unless they were discarded rather than moved, in which
+// case on what is left of the file.
+func (self *WorkingTreeDiffActions) revealSelectionInPaneItLandsIn(
+	pane types.DiffPaneContext, firstLineIdx int, actedSideRemains bool, otherSideHasChanges bool,
+) {
+	target := pane
+	if !actedSideRemains && otherSideHasChanges {
+		target = self.otherPane(pane)
+	}
+
+	revealSelectionAfterAction(self.c, pane, target, firstLineIdx)
+}
+
+// otherPane returns the main pane that isn't the given one.
+func (self *WorkingTreeDiffActions) otherPane(pane types.DiffPaneContext) types.DiffPaneContext {
+	if pane.GetKey() == self.c.Contexts().Normal.GetKey() {
+		return self.c.Contexts().NormalSecondary
+	}
+	return self.c.Contexts().Normal
+}
+
+// anyFileHasChangesOnSide reports whether any file under the selected node, other than
+// the ones named by except, has changes on the given side of the index, as the model
+// has them. The model is right about any file the action didn't touch; the ones it did
+// touch report for themselves, their entry not being right until the refresh lands.
+func (self *WorkingTreeDiffActions) anyFileHasChangesOnSide(except *set.Set[string], staged bool) bool {
+	node := self.context().GetSelected()
+	if node == nil {
+		return false
+	}
+
+	found := false
+	_ = node.ForEachFile(func(file *models.File) error {
+		if except.Includes(file.GetPath()) {
+			return nil
+		}
+		if (staged && file.HasStagedChanges) || (!staged && file.HasUnstagedChanges) {
+			found = true
+		}
+		return nil
+	})
+	return found
+}

pkg/gui/gui_common.go

@@ -185,6 +185,10 @@ func (self *guiCommon) GetViewBufferManagerForView(view *gocui.View) *tasks.View
 	return self.gui.getViewBufferManagerForView(view)
 }
 
+func (self *guiCommon) GetOrCreateViewBufferManagerForView(view *gocui.View) *tasks.ViewBufferManager {
+	return self.gui.getManager(view)
+}
+
 func (self *guiCommon) ReadLinesToFillView(view *gocui.View) {
 	self.gui.readLinesToFillView(view)
 }

pkg/gui/types/common.go

@@ -72,6 +72,10 @@ type IGuiCommon interface {
 	// return the view buffer manager for the given view, or nil if it doesn't have one
 	GetViewBufferManagerForView(view *gocui.View) *tasks.ViewBufferManager
 
+	// return the view buffer manager for the given view, making one if the view has
+	// never rendered anything, for saying something about a render still to come
+	GetOrCreateViewBufferManagerForView(view *gocui.View) *tasks.ViewBufferManager
+
 	// read enough lines into the given view's buffer to fill it at its current
 	// scroll position, plus some read-ahead for smooth scrolling
 	ReadLinesToFillView(view *gocui.View)

pkg/integration/tests/main_view/focus_follows_staged_side.go

@@ -0,0 +1,54 @@
+package main_view
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var FocusFollowsStagedSide = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Unstaging from a fully staged file leaves the focus on the staged side, which keeps its pane",
+	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")
+
+		// Two staged additions and nothing unstaged, so only the pane the staged side
+		// lives in is shown.
+		shell.UpdateFileAndAdd("file1", "one\nSTAGED1\ntwo\nthree\nfour\nfive\nsix\nseven\nSTAGED2\neight\nnine\nten\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Lines(
+				Contains("M  file1").IsSelected(),
+			).
+			Press(keys.Universal.FocusMainView)
+
+		t.Views().Main().IsInvisible()
+		t.Views().Secondary().
+			IsFocused().
+			Title(Equals("Staged changes")).
+			SelectedLines(
+				Contains("+STAGED1"),
+			).
+			PressPrimaryAction()
+
+		// The line taken out of the index turns up in the pane that has just appeared
+		// above, and the work carries on where it was, on the next staged change.
+		t.Views().Files().Lines(
+			Contains("MM file1"),
+		)
+		t.Views().Main().
+			IsVisible().
+			Content(Contains("+STAGED1"))
+		t.Views().Secondary().
+			IsFocused().
+			SelectedLines(
+				Contains("+STAGED2"),
+			)
+	},
+})

pkg/integration/tests/main_view/focus_follows_when_pane_goes.go

@@ -0,0 +1,47 @@
+package main_view
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var FocusFollowsWhenPaneGoes = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Staging the last unstaged change takes the upper pane away, so the focus follows the lines into the lower one",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(cfg *config.AppConfig) {
+		cfg.GetUserConfig().Gui.UseHunkModeInStagingView = false
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.CreateFileAndAdd("file1", "one\ntwo\nthree\n")
+		shell.Commit("one")
+
+		shell.UpdateFile("file1", "one\nADDED\ntwo\nthree\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Press(keys.Universal.FocusMainView)
+
+		t.Views().Main().
+			IsFocused().
+			Title(Equals("Unstaged changes")).
+			SelectedLines(
+				Contains("+ADDED"),
+			).
+			PressPrimaryAction()
+
+		// Nothing is unstaged any more, so that pane is gone and the line is in the one
+		// below, where the focus and the selection now are.
+		t.Views().Files().Lines(
+			Contains("M  file1"),
+		)
+		t.Views().Main().IsInvisible()
+		t.Views().Secondary().
+			IsFocused().
+			Title(Equals("Staged changes")).
+			SelectedLines(
+				Contains("+ADDED"),
+			)
+	},
+})

pkg/integration/tests/main_view/focus_returns_when_split_collapses.go

@@ -0,0 +1,56 @@
+package main_view
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var FocusReturnsWhenSplitCollapses = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Unstaging the last staged change from the secondary pane brings the focus back to the main one",
+	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")
+
+		// One change on each side, so the diff is split.
+		shell.UpdateFileAndAdd("file1", "one\nSTAGED\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten\n")
+		shell.UpdateFile("file1", "one\nSTAGED\ntwo\nthree\nfour\nfive\nsix\nseven\nUNSTAGED\neight\nnine\nten\n")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Files().
+			IsFocused().
+			Lines(
+				Contains("MM file1").IsSelected(),
+			).
+			Press(keys.Universal.FocusMainView)
+
+		t.Views().Main().
+			IsFocused().
+			Press(keys.Universal.TogglePanel)
+
+		t.Views().Secondary().
+			IsFocused().
+			SelectedLines(
+				Contains("+STAGED"),
+			).
+			PressPrimaryAction()
+
+		// With nothing staged left the diff isn't split any more, so the pane that was
+		// showing the staged side is gone — and the focus is back on the main one, where
+		// the change just taken out of the index now is.
+		t.Views().Files().Lines(
+			Contains(" M file1"),
+		)
+		t.Views().Main().
+			IsFocused().
+			Content(Contains("+STAGED")).
+			Content(Contains("+UNSTAGED")).
+			SelectedLines(
+				Contains("+STAGED"),
+			)
+	},
+})

pkg/integration/tests/test_list.go

@@ -376,7 +376,10 @@ var tests = []*components.IntegrationTest{
 	main_view.FileNavigationScrollsToTheTop,
 	main_view.FocusFollowsAPaneEmptiedFromOutside,
 	main_view.FocusFollowsIntoAPaneTakingOver,
+	main_view.FocusFollowsStagedSide,
+	main_view.FocusFollowsWhenPaneGoes,
 	main_view.FocusLeavesAnAlwaysSplitEmptyPane,
+	main_view.FocusReturnsWhenSplitCollapses,
 	main_view.HideSelectionWhenChangesVanish,
 	main_view.KeepAWrappedLineCoveredAcrossARerender,
 	main_view.KeepBothHalvesOfAChangeSelected,