WIP FocusedMainViewSnapshot approach

pkg/gui/context/patch_explorer_context.go

@@ -20,6 +20,11 @@ type PatchExplorerContext struct {
 	// true if we're inside the OnSelectItem callback; in that case we don't want to update the
 	// search result index.
 	inOnSelectItemCallback bool
+
+	// Set when this patch explorer was entered from a focused main view, so that
+	// escaping returns there; nil for the normal flow. See
+	// types.FocusedMainViewSnapshot.
+	focusedMainViewSnapshot *types.FocusedMainViewSnapshot
 }
 
 var (
@@ -60,6 +65,14 @@ func NewPatchExplorerContext(
 
 func (self *PatchExplorerContext) IsPatchExplorerContext() {}
 
+func (self *PatchExplorerContext) GetFocusedMainViewSnapshot() *types.FocusedMainViewSnapshot {
+	return self.focusedMainViewSnapshot
+}
+
+func (self *PatchExplorerContext) SetFocusedMainViewSnapshot(snapshot *types.FocusedMainViewSnapshot) {
+	self.focusedMainViewSnapshot = snapshot
+}
+
 func (self *PatchExplorerContext) GetState() *patch_exploring.State {
 	return self.state
 }

pkg/gui/controllers/commits_files_controller.go

@@ -548,6 +548,9 @@ func (self *CommitFilesController) expandAll() error {
 
 func (self *CommitFilesController) GetOnClickFocusedMainView() func(mainViewName string, clickedLineIdx int) error {
 	return func(mainViewName string, clickedLineIdx int) error {
+		// Capture before any mutation below that might re-render the main view.
+		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context(), clickedLineIdx)
+
 		clickedFile, line, ok := self.c.Helpers().Staging.GetFileAndLineForClickedDiffLine(mainViewName, clickedLineIdx)
 		if !ok {
 			line = -1
@@ -576,9 +579,8 @@ func (self *CommitFilesController) GetOnClickFocusedMainView() func(mainViewName
 			}
 		}
 
-		// Entered from the commit files panel's own focused main view, so escape
-		// should just pop back to it; no special escape context needed.
-		return self.c.Helpers().CommitFiles.EnterCommitFile(node, nil, types.OnFocusOpts{ClickedWindowName: "main", ClickedViewLineIdx: line, ClickedViewRealLineIdx: line})
+		// Entered from the focused main view, so escaping returns there.
+		return self.c.Helpers().CommitFiles.EnterCommitFile(node, snapshot, types.OnFocusOpts{ClickedWindowName: "main", ClickedViewLineIdx: line, ClickedViewRealLineIdx: line})
 	}
 }
 

pkg/gui/controllers/files_controller.go

@@ -409,6 +409,9 @@ func (self *FilesController) GetOnDoubleClick() func() error {
 
 func (self *FilesController) GetOnClickFocusedMainView() func(mainViewName string, clickedLineIdx int) error {
 	return func(mainViewName string, clickedLineIdx int) error {
+		// Capture before any mutation below that might re-render the main view.
+		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context(), clickedLineIdx)
+
 		clickedFile, line, ok := self.c.Helpers().Staging.GetFileAndLineForClickedDiffLine(mainViewName, clickedLineIdx)
 		if !ok {
 			line = -1
@@ -436,7 +439,7 @@ func (self *FilesController) GetOnClickFocusedMainView() func(mainViewName strin
 			}
 		}
 
-		return self.EnterFile(types.OnFocusOpts{ClickedWindowName: mainViewName, ClickedViewLineIdx: line, ClickedViewRealLineIdx: line})
+		return self.EnterFile(snapshot, types.OnFocusOpts{ClickedWindowName: mainViewName, ClickedViewLineIdx: line, ClickedViewRealLineIdx: line})
 	}
 }
 
@@ -724,7 +727,7 @@ func (self *FilesController) getSelectedFile() *models.File {
 }
 
 func (self *FilesController) enter() error {
-	return self.EnterFile(types.OnFocusOpts{ClickedWindowName: "", ClickedViewLineIdx: -1, ClickedViewRealLineIdx: -1})
+	return self.EnterFile(nil, types.OnFocusOpts{ClickedWindowName: "", ClickedViewLineIdx: -1, ClickedViewRealLineIdx: -1})
 }
 
 func (self *FilesController) collapseAll() error {
@@ -743,7 +746,11 @@ func (self *FilesController) expandAll() error {
 	return nil
 }
 
-func (self *FilesController) EnterFile(opts types.OnFocusOpts) error {
+// focusedMainViewSnapshot records the focused main view to return to when
+// escaping the staging view, for the case where we're entering it straight from
+// there; it's nil for the normal flow that goes through the files panel. See
+// types.FocusedMainViewSnapshot.
+func (self *FilesController) EnterFile(focusedMainViewSnapshot *types.FocusedMainViewSnapshot, opts types.OnFocusOpts) error {
 	node := self.context().GetSelected()
 	if node == nil {
 		return nil
@@ -770,6 +777,9 @@ func (self *FilesController) EnterFile(opts types.OnFocusOpts) error {
 	}
 
 	context := lo.Ternary(opts.ClickedWindowName == "secondary", self.c.Contexts().StagingSecondary, self.c.Contexts().Staging)
+	// Set on every entry (so it can't leak from a previous main-view entry into a
+	// subsequent normal one), right as we push the staging view.
+	context.SetFocusedMainViewSnapshot(focusedMainViewSnapshot)
 	self.c.Context().Push(context, opts)
 	self.c.Helpers().PatchBuilding.ShowHunkStagingHint()
 
@@ -1553,7 +1563,7 @@ func (self *FilesController) handleStashSave(stashFunc func(message string) erro
 }
 
 func (self *FilesController) onClickMain(opts gocui.ViewMouseBindingOpts) error {
-	return self.EnterFile(types.OnFocusOpts{ClickedWindowName: "main", ClickedViewLineIdx: opts.Y})
+	return self.EnterFile(nil, types.OnFocusOpts{ClickedWindowName: "main", ClickedViewLineIdx: opts.Y})
 }
 
 func (self *FilesController) fetch() error {

pkg/gui/controllers/helpers/commit_files_helper.go

@@ -21,11 +21,11 @@ func NewCommitFilesHelper(c *HelperCommon, patchBuildingHelper *PatchBuildingHel
 	}
 }
 
-// escapeContext is the side panel that escaping the patch builder should return
-// to, for the case where we're entering it straight from a focused main view;
-// it's nil for the normal flow that goes through the commit files panel. See
-// PatchBuildingHelper.escapeContext.
-func (self *CommitFilesHelper) EnterCommitFile(node *filetree.CommitFileNode, escapeContext types.Context, opts types.OnFocusOpts) error {
+// focusedMainViewSnapshot records the focused main view to return to when
+// escaping the patch builder, for the case where we're entering it straight from
+// there; it's nil for the normal flow that goes through the commit files panel.
+// See types.FocusedMainViewSnapshot.
+func (self *CommitFilesHelper) EnterCommitFile(node *filetree.CommitFileNode, focusedMainViewSnapshot *types.FocusedMainViewSnapshot, opts types.OnFocusOpts) error {
 	if node.File == nil {
 		self.handleToggleCommitFileDirCollapsed(node)
 		return nil
@@ -55,7 +55,7 @@ func (self *CommitFilesHelper) EnterCommitFile(node *filetree.CommitFileNode, es
 			// Set on every entry (so it can't leak from a previous main-view
 			// entry into a subsequent normal one), right as we push the patch
 			// builder.
-			self.patchBuildingHelper.escapeContext = escapeContext
+			self.c.Contexts().CustomPatchBuilder.SetFocusedMainViewSnapshot(focusedMainViewSnapshot)
 
 			self.c.Context().Push(self.c.Contexts().CustomPatchBuilder, opts)
 			self.patchBuildingHelper.ShowHunkStagingHint()

pkg/gui/controllers/helpers/patch_building_helper.go

@@ -10,14 +10,6 @@ import (
 
 type PatchBuildingHelper struct {
 	c *HelperCommon
-
-	// When patch building is entered straight from a focused main view (rather
-	// than from the commit files panel), this records the side panel to return
-	// to on escape, so that we skip the commit files panel we never really
-	// visited. It is nil for the normal flow, where escape just pops back to the
-	// commit files panel. Set on every entry into patch building (see
-	// CommitFilesHelper.EnterCommitFile) so it can't leak between flows.
-	escapeContext types.Context
 }
 
 func NewPatchBuildingHelper(
@@ -40,17 +32,47 @@ func (self *PatchBuildingHelper) ShowHunkStagingHint() {
 	}
 }
 
-// takes us from the patch building panel back to the commit files panel, or
-// straight back to the side panel if we entered patch building from a focused
-// main view (see escapeContext)
+// takes us from the patch building panel back to the commit files panel, or to
+// the focused main view if that's where we entered it from
 func (self *PatchBuildingHelper) Escape() {
-	if self.escapeContext != nil {
-		escapeContext := self.escapeContext
-		self.escapeContext = nil
-		self.c.Context().Push(escapeContext, types.OnFocusOpts{})
-	} else {
-		self.c.Context().Pop()
+	EscapeFromPatchExplorer(self.c, self.c.Contexts().CustomPatchBuilder)
+}
+
+// EscapeFromPatchExplorer returns from a patch explorer context (staging or
+// patch building). If we entered it from a focused main view, we go back to
+// where we came from (re-rendering the side panel's content into the main view,
+// like the plain escape does), then focus the main view and restore its scroll
+// position and selection. Otherwise we just pop to the side panel.
+func EscapeFromPatchExplorer(c *HelperCommon, context types.IPatchExplorerContext) {
+	snapshot := context.GetFocusedMainViewSnapshot()
+	if snapshot == nil {
+		c.Context().Pop()
+		return
 	}
+
+	context.SetFocusedMainViewSnapshot(nil)
+
+	// Restore the side panel's selection before we render it, so it shows the
+	// same content the main view had (diving into staging can change it, e.g.
+	// from a directory to a file in the files panel).
+	if listContext, ok := snapshot.SidePanel.(types.IListContext); ok && snapshot.SidePanelSelectedLineIdx >= 0 {
+		listContext.GetList().SetSelectedLineIdx(snapshot.SidePanelSelectedLineIdx)
+	}
+
+	// Land on the side panel first (this re-renders the original content into the
+	// main view), then focus the main view on top of it.
+	c.Context().Push(snapshot.SidePanel, types.OnFocusOpts{})
+	c.Context().Push(snapshot.MainView, types.OnFocusOpts{})
+
+	// Restore the scroll position and selection on the next UI tick.
+	view := snapshot.MainView.GetView()
+	c.OnUIThread(func() error {
+		view.SetOrigin(view.OriginX(), snapshot.OriginY)
+		view.FocusPoint(0, snapshot.SelectedLineIdx, false)
+		view.Highlight = true
+		view.HighlightInactive = false
+		return nil
+	})
 }
 
 // kills the custom patch and returns us back to the commit files panel if needed

pkg/gui/controllers/main_view_controller.go

@@ -192,6 +192,37 @@ func showSelectionAtLine(view *gocui.View, lineIdx int) {
 	view.FocusPoint(0, lineIdx, false)
 }
 
+// focusedMainViewContextForViewName maps a focused main view's view name (as
+// passed to GetOnClickFocusedMainView) to its context.
+func focusedMainViewContextForViewName(c *ControllerCommon, viewName string) types.Context {
+	if viewName == c.Contexts().NormalSecondary.GetViewName() {
+		return c.Contexts().NormalSecondary
+	}
+	return c.Contexts().Normal
+}
+
+// focusedMainViewSnapshot captures where a focused main view is (scroll +
+// selected line) when diving into a patch explorer from it, so escaping can
+// return there with the main view focused. sidePanel is the panel to land on
+// first (which re-renders the content); for commits/stash it's the originating
+// panel, skipping the commit files panel we pass through. selectedLineIdx is the
+// view line that was selected in the focused main view. Call this before any
+// mutation that might re-render the main view.
+func focusedMainViewSnapshot(c *ControllerCommon, mainViewName string, sidePanel types.Context, selectedLineIdx int) *types.FocusedMainViewSnapshot {
+	mainView := focusedMainViewContextForViewName(c, mainViewName)
+	sidePanelSelectedLineIdx := -1
+	if listContext, ok := sidePanel.(types.IListContext); ok {
+		sidePanelSelectedLineIdx = listContext.GetList().GetSelectedLineIdx()
+	}
+	return &types.FocusedMainViewSnapshot{
+		SidePanel:                sidePanel,
+		SidePanelSelectedLineIdx: sidePanelSelectedLineIdx,
+		MainView:                 mainView,
+		OriginY:                  mainView.GetView().OriginY(),
+		SelectedLineIdx:          selectedLineIdx,
+	}
+}
+
 func (self *MainViewController) editLine() error {
 	if !self.context.GetView().Highlight {
 		return nil

pkg/gui/controllers/staging_controller.go

@@ -7,6 +7,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/patch"
 	"github.com/jesseduffield/lazygit/pkg/gocui"
+	"github.com/jesseduffield/lazygit/pkg/gui/controllers/helpers"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 )
 
@@ -175,7 +176,7 @@ func (self *StagingController) Escape() error {
 		return nil
 	}
 
-	self.c.Context().Pop()
+	helpers.EscapeFromPatchExplorer(self.c.HelperCommon, self.context)
 	return nil
 }
 

pkg/gui/controllers/switch_to_diff_files_controller.go

@@ -59,6 +59,11 @@ func (self *SwitchToDiffFilesController) GetOnClickFocusedMainView() func(mainVi
 			return nil
 		}
 
+		// Capture before self.enter() pushes the commit files panel, which
+		// re-renders the main view. We escape "all the way out" to this side
+		// panel (skipping the commit files panel), then focus the main view.
+		snapshot := focusedMainViewSnapshot(self.c, mainViewName, self.context, clickedLineIdx)
+
 		if err := self.enter(); err != nil {
 			return err
 		}
@@ -83,10 +88,7 @@ func (self *SwitchToDiffFilesController) GetOnClickFocusedMainView() func(mainVi
 		context.GetViewTrait().FocusPoint(
 			context.ModelIndexToViewIndex(idx), false)
 		node = context.GetSelected()
-		// We entered patch building straight from the focused main view, so
-		// escaping it should take us all the way back out to this side panel,
-		// skipping the commit files panel we never really visited.
-		return self.c.Helpers().CommitFiles.EnterCommitFile(node, self.context, types.OnFocusOpts{ClickedWindowName: "main", ClickedViewLineIdx: line, ClickedViewRealLineIdx: line})
+		return self.c.Helpers().CommitFiles.EnterCommitFile(node, snapshot, types.OnFocusOpts{ClickedWindowName: "main", ClickedViewLineIdx: line, ClickedViewRealLineIdx: line})
 	}
 }
 

pkg/gui/types/context.go

@@ -205,6 +205,34 @@ type IPatchExplorerContext interface {
 	NavigateTo(selectedLineIdx int)
 	GetMutex() *deadlock.Mutex
 	IsPatchExplorerContext() // used for type switch
+
+	// See FocusedMainViewSnapshot. Nil unless this patch explorer was entered
+	// from a focused main view.
+	GetFocusedMainViewSnapshot() *FocusedMainViewSnapshot
+	SetFocusedMainViewSnapshot(*FocusedMainViewSnapshot)
+}
+
+// FocusedMainViewSnapshot records where a focused main view was when we dived
+// into a patch explorer (staging or patch building) from it, so that escaping
+// returns us to the same place with the main view focused again. It is nil when
+// the patch explorer was entered the normal way (through a side panel), in which
+// case escape just pops to that side panel.
+type FocusedMainViewSnapshot struct {
+	// The side panel to land on first; pushing it re-renders the original
+	// content into the main view. For commits/stash this is the originating side
+	// panel (skipping the commit files panel we passed through), preserving the
+	// pre-existing "escape all the way out" behavior.
+	SidePanel Context
+	// The side panel's selected line, to restore before re-rendering it. Diving
+	// into staging can change the side panel's selection (e.g. from a directory
+	// to a file in the files panel); restoring it makes the main view show the
+	// same content again. -1 if the side panel isn't a list.
+	SidePanelSelectedLineIdx int
+	// The focused main view context to focus afterwards.
+	MainView Context
+	// The scroll position and selected line to restore in the main view.
+	OriginY         int
+	SelectedLineIdx int
 }
 
 type IViewTrait interface {