Land in the same panel when switching to a worktreecopper

11 files changed, 50 insertions, 16 deletions

diff --git a/pkg/gui/context.go b/pkg/gui/context.go
index 38e1212c7..90bb78768 100644
--- a/pkg/gui/context.go
+++ b/pkg/gui/context.go
@@ -376,3 +376,16 @@ func (self *ContextMgr) AllPatchExplorer() []types.IPatchExplorerContext {
 
 	return listContexts
 }
+
+func (self *ContextMgr) ContextForKey(key types.ContextKey) types.Context {
+	self.RLock()
+	defer self.RUnlock()
+
+	for _, context := range self.allContexts.Flatten() {
+		if context.GetKey() == key {
+			return context
+		}
+	}
+
+	return nil
+}
diff --git a/pkg/gui/context/context.go b/pkg/gui/context/context.go
index 12fc285ae..ecb0e6da4 100644
--- a/pkg/gui/context/context.go
+++ b/pkg/gui/context/context.go
@@ -5,6 +5,9 @@ import (
 )
 
 const (
+	// used as a nil value when passing a context key as an arg
+	NO_CONTEXT types.ContextKey = "none"
+
 	GLOBAL_CONTEXT_KEY                   types.ContextKey = "global"
 	STATUS_CONTEXT_KEY                   types.ContextKey = "status"
 	SNAKE_CONTEXT_KEY                    types.ContextKey = "snake"
diff --git a/pkg/gui/controllers/branches_controller.go b/pkg/gui/controllers/branches_controller.go
index 5ce8d88fb..436fbb211 100644
--- a/pkg/gui/controllers/branches_controller.go
+++ b/pkg/gui/controllers/branches_controller.go
@@ -228,7 +228,7 @@ func (self *BranchesController) promptToCheckoutWorktree(worktree *models.Worktr
 		Title:  "Switch to worktree",
 		Prompt: fmt.Sprintf("This branch is checked out by worktree %s. Do you want to switch to that worktree?", worktree.Name()),
 		HandleConfirm: func() error {
-			return self.c.Helpers().Worktree.Switch(worktree)
+			return self.c.Helpers().Worktree.Switch(worktree, context.LOCAL_BRANCHES_CONTEXT_KEY)
 		},
 	})
 }
diff --git a/pkg/gui/controllers/helpers/repos_helper.go b/pkg/gui/controllers/helpers/repos_helper.go
index 013580ff2..942cceef8 100644
--- a/pkg/gui/controllers/helpers/repos_helper.go
+++ b/pkg/gui/controllers/helpers/repos_helper.go
@@ -12,13 +12,14 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/env"
+	"github.com/jesseduffield/lazygit/pkg/gui/context"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation/icons"
 	"github.com/jesseduffield/lazygit/pkg/gui/style"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
-type onNewRepoFn func(startArgs appTypes.StartArgs, reuseState bool) error
+type onNewRepoFn func(startArgs appTypes.StartArgs, reuseState bool, contextKey types.ContextKey) error
 
 // helps switch back and forth between repos
 type ReposHelper struct {
@@ -46,7 +47,7 @@ func (self *ReposHelper) EnterSubmodule(submodule *models.SubmoduleConfig) error
 	}
 	self.c.State().GetRepoPathStack().Push(wd)
 
-	return self.DispatchSwitchToRepo(submodule.Path, true)
+	return self.DispatchSwitchToRepo(submodule.Path, true, context.NO_CONTEXT)
 }
 
 func (self *ReposHelper) getCurrentBranch(path string) string {
@@ -129,7 +130,7 @@ func (self *ReposHelper) CreateRecentReposMenu() error {
 				// if we were in a submodule, we want to forget about that stack of repos
 				// so that hitting escape in the new repo does nothing
 				self.c.State().GetRepoPathStack().Clear()
-				return self.DispatchSwitchToRepo(path, false)
+				return self.DispatchSwitchToRepo(path, false, context.NO_CONTEXT)
 			},
 		}
 	})
@@ -137,11 +138,11 @@ func (self *ReposHelper) CreateRecentReposMenu() error {
 	return self.c.Menu(types.CreateMenuOptions{Title: self.c.Tr.RecentRepos, Items: menuItems})
 }
 
-func (self *ReposHelper) DispatchSwitchToRepo(path string, reuse bool) error {
-	return self.DispatchSwitchTo(path, reuse, self.c.Tr.ErrRepositoryMovedOrDeleted)
+func (self *ReposHelper) DispatchSwitchToRepo(path string, reuse bool, contextKey types.ContextKey) error {
+	return self.DispatchSwitchTo(path, reuse, self.c.Tr.ErrRepositoryMovedOrDeleted, contextKey)
 }
 
-func (self *ReposHelper) DispatchSwitchTo(path string, reuse bool, errMsg string) error {
+func (self *ReposHelper) DispatchSwitchTo(path string, reuse bool, errMsg string, contextKey types.ContextKey) error {
 	env.UnsetGitDirEnvs()
 	originalPath, err := os.Getwd()
 	if err != nil {
@@ -175,5 +176,5 @@ func (self *ReposHelper) DispatchSwitchTo(path string, reuse bool, errMsg string
 	self.c.Mutexes().RefreshingFilesMutex.Lock()
 	defer self.c.Mutexes().RefreshingFilesMutex.Unlock()
 
-	return self.onNewRepo(appTypes.StartArgs{}, reuse)
+	return self.onNewRepo(appTypes.StartArgs{}, reuse, contextKey)
 }
diff --git a/pkg/gui/controllers/helpers/worktree_helper.go b/pkg/gui/controllers/helpers/worktree_helper.go
index 888838370..b43fb7d69 100644
--- a/pkg/gui/controllers/helpers/worktree_helper.go
+++ b/pkg/gui/controllers/helpers/worktree_helper.go
@@ -78,7 +78,7 @@ func (self *WorktreeHelper) NewWorktree() error {
 	})
 }
 
-func (self *WorktreeHelper) Switch(worktree *models.Worktree) error {
+func (self *WorktreeHelper) Switch(worktree *models.Worktree, contextKey types.ContextKey) error {
 	if self.c.Git().Worktree.IsCurrentWorktree(worktree) {
 		return self.c.ErrorMsg(self.c.Tr.AlreadyInWorktree)
 	}
@@ -89,5 +89,5 @@ func (self *WorktreeHelper) Switch(worktree *models.Worktree) error {
 	// so that hitting escape in the new repo does nothing
 	self.c.State().GetRepoPathStack().Clear()
 
-	return self.reposHelper.DispatchSwitchTo(worktree.Path, true, self.c.Tr.ErrWorktreeMovedOrDeleted)
+	return self.reposHelper.DispatchSwitchTo(worktree.Path, true, self.c.Tr.ErrWorktreeMovedOrRemoved, contextKey)
 }
diff --git a/pkg/gui/controllers/quit_actions.go b/pkg/gui/controllers/quit_actions.go
index a163f66c8..a7301d5b1 100644
--- a/pkg/gui/controllers/quit_actions.go
+++ b/pkg/gui/controllers/quit_actions.go
@@ -2,6 +2,7 @@ package controllers
 
 import (
 	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/gui/context"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 )
 
@@ -77,7 +78,7 @@ func (self *QuitActions) Escape() error {
 
 	repoPathStack := self.c.State().GetRepoPathStack()
 	if !repoPathStack.IsEmpty() {
-		return self.c.Helpers().Repos.DispatchSwitchToRepo(repoPathStack.Pop(), true)
+		return self.c.Helpers().Repos.DispatchSwitchToRepo(repoPathStack.Pop(), true, context.NO_CONTEXT)
 	}
 
 	if self.c.UserConfig.QuitOnTopLevelReturn {
diff --git a/pkg/gui/controllers/worktrees_controller.go b/pkg/gui/controllers/worktrees_controller.go
index 3c5d5f381..2f5a35bd5 100644
--- a/pkg/gui/controllers/worktrees_controller.go
+++ b/pkg/gui/controllers/worktrees_controller.go
@@ -143,7 +143,7 @@ func (self *WorktreesController) GetOnClick() func() error {
 }
 
 func (self *WorktreesController) enter(worktree *models.Worktree) error {
-	return self.c.Helpers().Worktree.Switch(worktree)
+	return self.c.Helpers().Worktree.Switch(worktree, context.WORKTREES_CONTEXT_KEY)
 }
 
 func (self *WorktreesController) checkSelected(callback func(worktree *models.Worktree) error) func() error {
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index a0ff3747b..626240ed3 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -274,7 +274,7 @@ func (self *GuiRepoState) GetSplitMainPanel() bool {
 	return self.SplitMainPanel
 }
 
-func (gui *Gui) onNewRepo(startArgs appTypes.StartArgs, reuseState bool) error {
+func (gui *Gui) onNewRepo(startArgs appTypes.StartArgs, reuseState bool, contextKey types.ContextKey) error {
 	var err error
 	gui.git, err = commands.NewGitCommand(
 		gui.Common,
@@ -295,6 +295,17 @@ func (gui *Gui) onNewRepo(startArgs appTypes.StartArgs, reuseState bool) error {
 		return err
 	}
 
+	// if a context key has been given, push that instead, and set its index to 0
+	if contextKey != context.NO_CONTEXT {
+		contextToPush = gui.c.ContextForKey(contextKey)
+		// when we pass a list context, the expectation is that our cursor goes to the top,
+		// because e.g. with worktrees, we'll show the current worktree at the top of the list.
+		listContext, ok := contextToPush.(types.IListContext)
+		if ok {
+			listContext.GetList().SetSelectedLineIdx(0)
+		}
+	}
+
 	if err := gui.c.PushContext(contextToPush); err != nil {
 		return err
 	}
@@ -618,7 +629,7 @@ func (gui *Gui) Run(startArgs appTypes.StartArgs) error {
 	}
 
 	// onNewRepo must be called after g.SetManager because SetManager deletes keybindings
-	if err := gui.onNewRepo(startArgs, false); err != nil {
+	if err := gui.onNewRepo(startArgs, false, context.NO_CONTEXT); err != nil {
 		return err
 	}
 
diff --git a/pkg/gui/gui_common.go b/pkg/gui/gui_common.go
index c0d7bd460..91f24d448 100644
--- a/pkg/gui/gui_common.go
+++ b/pkg/gui/gui_common.go
@@ -76,6 +76,10 @@ func (self *guiCommon) Context() types.IContextMgr {
 	return self.gui.State.ContextMgr
 }
 
+func (self *guiCommon) ContextForKey(key types.ContextKey) types.Context {
+	return self.gui.State.ContextMgr.ContextForKey(key)
+}
+
 func (self *guiCommon) ActivateContext(context types.Context) error {
 	return self.gui.State.ContextMgr.ActivateContext(context, types.OnFocusOpts{})
 }
diff --git a/pkg/gui/types/common.go b/pkg/gui/types/common.go
index 4ffeb2791..e23354cbd 100644
--- a/pkg/gui/types/common.go
+++ b/pkg/gui/types/common.go
@@ -66,6 +66,7 @@ type IGuiCommon interface {
 	IsCurrentContext(Context) bool
 	// TODO: replace the above context-based methods with just using Context() e.g. replace PushContext() with Context().Push()
 	Context() IContextMgr
+	ContextForKey(key ContextKey) Context
 
 	ActivateContext(context Context) error
 
diff --git a/pkg/i18n/english.go b/pkg/i18n/english.go
index f4479c34a..da05da96b 100644
--- a/pkg/i18n/english.go
+++ b/pkg/i18n/english.go
@@ -480,7 +480,7 @@ type TranslationSet struct {
 	ErrCannotEditDirectory              string
 	ErrStageDirWithInlineMergeConflicts string
 	ErrRepositoryMovedOrDeleted         string
-	ErrWorktreeMovedOrDeleted           string
+	ErrWorktreeMovedOrRemoved           string
 	CommandLog                          string
 	ToggleShowCommandLog                string
 	FocusCommandLog                     string
@@ -1199,7 +1199,7 @@ func EnglishTranslationSet() TranslationSet {
 		ErrStageDirWithInlineMergeConflicts: "Cannot stage/unstage directory containing files with inline merge conflicts. Please fix up the merge conflicts first",
 		ErrRepositoryMovedOrDeleted:         "Cannot find repo. It might have been moved or deleted ¯\\_(ツ)_/¯",
 		CommandLog:                          "Command log",
-		ErrWorktreeMovedOrDeleted:           "Cannot find worktree. It might have been moved or deleted ¯\\_(ツ)_/¯",
+		ErrWorktreeMovedOrRemoved:           "Cannot find worktree. It might have been moved or removed ¯\\_(ツ)_/¯",
 		ToggleShowCommandLog:                "Toggle show/hide command log",
 		FocusCommandLog:                     "Focus command log",
 		CommandLogHeader:                    "You can hide/focus this panel by pressing '%s'\n",