Support creating worktrees from refs

6 files changed, 199 insertions, 17 deletions

diff --git a/pkg/gui/controllers.go b/pkg/gui/controllers.go
index c02d6bae2..893c59dbe 100644
--- a/pkg/gui/controllers.go
+++ b/pkg/gui/controllers.go
@@ -18,6 +18,9 @@ func (gui *Gui) Helpers() *helpers.Helpers {
 
 func (gui *Gui) resetHelpersAndControllers() {
 	helperCommon := gui.c
+	recordDirectoryHelper := helpers.NewRecordDirectoryHelper(helperCommon)
+	reposHelper := helpers.NewRecentReposHelper(helperCommon, recordDirectoryHelper, gui.onNewRepo)
+	worktreeHelper := helpers.NewWorktreeHelper(helperCommon, reposHelper)
 	refsHelper := helpers.NewRefsHelper(helperCommon)
 
 	rebaseHelper := helpers.NewMergeAndRebaseHelper(helperCommon, refsHelper)
@@ -41,12 +44,10 @@ func (gui *Gui) resetHelpersAndControllers() {
 
 	gpgHelper := helpers.NewGpgHelper(helperCommon)
 	viewHelper := helpers.NewViewHelper(helperCommon, gui.State.Contexts)
-	recordDirectoryHelper := helpers.NewRecordDirectoryHelper(helperCommon)
 	patchBuildingHelper := helpers.NewPatchBuildingHelper(helperCommon)
 	stagingHelper := helpers.NewStagingHelper(helperCommon)
 	mergeConflictsHelper := helpers.NewMergeConflictsHelper(helperCommon)
-	reposHelper := helpers.NewRecentReposHelper(helperCommon, recordDirectoryHelper, gui.onNewRepo)
-	worktreeHelper := helpers.NewWorktreeHelper(helperCommon, reposHelper)
+
 	refreshHelper := helpers.NewRefreshHelper(
 		helperCommon,
 		refsHelper,
@@ -241,6 +242,18 @@ func (gui *Gui) resetHelpersAndControllers() {
 		controllers.AttachControllers(context, controllers.NewBasicCommitsController(common, context))
 	}
 
+	for _, context := range []controllers.CanViewWorktreeOptions{
+		gui.State.Contexts.LocalCommits,
+		gui.State.Contexts.ReflogCommits,
+		gui.State.Contexts.SubCommits,
+		gui.State.Contexts.Stash,
+		gui.State.Contexts.Branches,
+		gui.State.Contexts.RemoteBranches,
+		gui.State.Contexts.Tags,
+	} {
+		controllers.AttachControllers(context, controllers.NewWorktreeOptionsController(common, context))
+	}
+
 	controllers.AttachControllers(gui.State.Contexts.ReflogCommits,
 		reflogCommitsController,
 	)
diff --git a/pkg/gui/controllers/branches_controller.go b/pkg/gui/controllers/branches_controller.go
index 1935b8d1c..3ead83f35 100644
--- a/pkg/gui/controllers/branches_controller.go
+++ b/pkg/gui/controllers/branches_controller.go
@@ -204,7 +204,7 @@ func (self *BranchesController) press(selectedBranch *models.Branch) error {
 
 	if selectedBranch.CheckedOutByOtherWorktree {
 		worktreeForRef, ok := self.worktreeForBranch(selectedBranch)
-		if ok && !self.c.Git().Worktree.IsCurrentWorktree(worktreeForRef) {
+		if ok && !self.c.Git().Worktree.IsCurrentWorktree(worktreeForRef.Path) {
 			return self.promptToCheckoutWorktree(worktreeForRef)
 		}
 	}
@@ -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, context.LOCAL_BRANCHES_CONTEXT_KEY)
+			return self.c.Helpers().Worktree.Switch(worktree.Path, context.LOCAL_BRANCHES_CONTEXT_KEY)
 		},
 	})
 }
@@ -346,7 +346,7 @@ func (self *BranchesController) promptWorktreeBranchDelete(selectedBranch *model
 			{
 				Label: "Switch to worktree",
 				OnPress: func() error {
-					return self.c.Helpers().Worktree.Switch(worktree, context.LOCAL_BRANCHES_CONTEXT_KEY)
+					return self.c.Helpers().Worktree.Switch(worktree.Path, context.LOCAL_BRANCHES_CONTEXT_KEY)
 				},
 			},
 			{
diff --git a/pkg/gui/controllers/helpers/worktree_helper.go b/pkg/gui/controllers/helpers/worktree_helper.go
index 321ec3ed0..48c13236e 100644
--- a/pkg/gui/controllers/helpers/worktree_helper.go
+++ b/pkg/gui/controllers/helpers/worktree_helper.go
@@ -9,7 +9,9 @@ import (
 	"strings"
 
 	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/jesseduffield/lazygit/pkg/gui/context"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/jesseduffield/lazygit/pkg/utils"
 )
@@ -66,10 +68,14 @@ func (self *WorktreeHelper) NewWorktree() error {
 		HandleConfirm: func(path string) error {
 			return self.c.Prompt(types.PromptOpts{
 				Title: self.c.Tr.NewWorktreeBranch,
-				HandleConfirm: func(committish string) error {
+				// TODO: suggestions
+				HandleConfirm: func(base string) error {
 					return self.c.WithWaitingStatus(self.c.Tr.AddingWorktree, func(gocui.Task) error {
 						self.c.LogAction(self.c.Tr.Actions.AddWorktree)
-						if err := self.c.Git().Worktree.New(sanitizedBranchName(path), committish); err != nil {
+						if err := self.c.Git().Worktree.New(git_commands.NewWorktreeOpts{
+							Path: path,
+							Base: base,
+						}); err != nil {
 							return err
 						}
 						return self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC, Scope: []types.RefreshableView{types.WORKTREES, types.BRANCHES, types.FILES}})
@@ -80,14 +86,70 @@ func (self *WorktreeHelper) NewWorktree() error {
 	})
 }
 
-func (self *WorktreeHelper) Switch(worktree *models.Worktree, contextKey types.ContextKey) error {
-	if self.c.Git().Worktree.IsCurrentWorktree(worktree) {
+func (self *WorktreeHelper) NewWorktreeCheckout(base string, isBranch bool, detached bool) error {
+	opts := git_commands.NewWorktreeOpts{
+		Base:   base,
+		Detach: detached,
+	}
+
+	f := func() error {
+		return self.c.WithWaitingStatus(self.c.Tr.AddingWorktree, func(gocui.Task) error {
+			self.c.LogAction(self.c.Tr.Actions.AddWorktree)
+			if err := self.c.Git().Worktree.New(opts); err != nil {
+				return err
+			}
+			return self.Switch(opts.Path, context.LOCAL_BRANCHES_CONTEXT_KEY)
+		})
+	}
+
+	return self.c.Prompt(types.PromptOpts{
+		Title: self.c.Tr.NewWorktreePath,
+		HandleConfirm: func(path string) error {
+			opts.Path = path
+
+			if detached {
+				return f()
+			}
+
+			if isBranch {
+				// prompt for the new branch name where a blank means we just check out the branch
+				return self.c.Prompt(types.PromptOpts{
+					Title: fmt.Sprintf("New branch name (leave blank to checkout %s)", base),
+					// TODO: suggestions
+					HandleConfirm: func(branchName string) error {
+						opts.Branch = branchName
+
+						return f()
+					},
+				})
+			} else {
+				// prompt for the new branch name where a blank means we just check out the branch
+				return self.c.Prompt(types.PromptOpts{
+					Title: "New branch name",
+					// TODO: suggestions
+					HandleConfirm: func(branchName string) error {
+						if branchName == "" {
+							return self.c.ErrorMsg("Branch name cannot be blank")
+						}
+
+						opts.Branch = branchName
+
+						return f()
+					},
+				})
+			}
+		},
+	})
+}
+
+func (self *WorktreeHelper) Switch(path string, contextKey types.ContextKey) error {
+	if self.c.Git().Worktree.IsCurrentWorktree(path) {
 		return self.c.ErrorMsg(self.c.Tr.AlreadyInWorktree)
 	}
 
 	self.c.LogAction(self.c.Tr.SwitchToWorktree)
 
-	return self.reposHelper.DispatchSwitchTo(worktree.Path, true, self.c.Tr.ErrWorktreeMovedOrRemoved, contextKey)
+	return self.reposHelper.DispatchSwitchTo(path, true, self.c.Tr.ErrWorktreeMovedOrRemoved, contextKey)
 }
 
 func (self *WorktreeHelper) Remove(worktree *models.Worktree, force bool) error {
@@ -139,3 +201,51 @@ func (self *WorktreeHelper) Detach(worktree *models.Worktree) error {
 		return self.c.Refresh(types.RefreshOptions{Mode: types.ASYNC, Scope: []types.RefreshableView{types.WORKTREES, types.BRANCHES, types.FILES}})
 	})
 }
+
+func (self *WorktreeHelper) ViewWorktreeOptions(context types.IListContext, ref string) error {
+	if context == self.c.Contexts().Branches {
+		return self.ViewBranchWorktreeOptions(ref)
+	}
+
+	return self.ViewRefWorktreeOptions(ref)
+}
+
+func (self *WorktreeHelper) ViewBranchWorktreeOptions(branchName string) error {
+	return self.c.Menu(types.CreateMenuOptions{
+		Title: self.c.Tr.WorktreeTitle,
+		Items: []*types.MenuItem{
+			{
+				LabelColumns: []string{"Create new worktree from branch"},
+				OnPress: func() error {
+					return self.NewWorktreeCheckout(branchName, true, false)
+				},
+			},
+			{
+				LabelColumns: []string{"Create new worktree from branch (detached)"},
+				OnPress: func() error {
+					return self.NewWorktreeCheckout(branchName, true, true)
+				},
+			},
+		},
+	})
+}
+
+func (self *WorktreeHelper) ViewRefWorktreeOptions(ref string) error {
+	return self.c.Menu(types.CreateMenuOptions{
+		Title: self.c.Tr.WorktreeTitle,
+		Items: []*types.MenuItem{
+			{
+				LabelColumns: []string{"Create new worktree from ref"},
+				OnPress: func() error {
+					return self.NewWorktreeCheckout(ref, false, false)
+				},
+			},
+			{
+				LabelColumns: []string{"Create new worktree from ref (detached)"},
+				OnPress: func() error {
+					return self.NewWorktreeCheckout(ref, false, true)
+				},
+			},
+		},
+	})
+}
diff --git a/pkg/gui/controllers/worktree_options_controller.go b/pkg/gui/controllers/worktree_options_controller.go
new file mode 100644
index 000000000..8c2c0bbb0
--- /dev/null
+++ b/pkg/gui/controllers/worktree_options_controller.go
@@ -0,0 +1,59 @@
+package controllers
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/gui/types"
+)
+
+// This controller is for all contexts that have items you can create a worktree from
+
+var _ types.IController = &WorktreeOptionsController{}
+
+type CanViewWorktreeOptions interface {
+	types.IListContext
+}
+
+type WorktreeOptionsController struct {
+	baseController
+	c       *ControllerCommon
+	context CanViewWorktreeOptions
+}
+
+func NewWorktreeOptionsController(controllerCommon *ControllerCommon, context CanViewWorktreeOptions) *WorktreeOptionsController {
+	return &WorktreeOptionsController{
+		baseController: baseController{},
+		c:              controllerCommon,
+		context:        context,
+	}
+}
+
+func (self *WorktreeOptionsController) GetKeybindings(opts types.KeybindingsOpts) []*types.Binding {
+	bindings := []*types.Binding{
+		{
+			Key:         opts.GetKey(opts.Config.Worktrees.ViewWorktreeOptions),
+			Handler:     self.checkSelected(self.viewWorktreeOptions),
+			Description: self.c.Tr.ViewWorktreeOptions,
+			OpensMenu:   true,
+		},
+	}
+
+	return bindings
+}
+
+func (self *WorktreeOptionsController) checkSelected(callback func(string) error) func() error {
+	return func() error {
+		ref := self.context.GetSelectedItemId()
+		if ref == "" {
+			return nil
+		}
+
+		return callback(ref)
+	}
+}
+
+func (self *WorktreeOptionsController) Context() types.Context {
+	return self.context
+}
+
+func (self *WorktreeOptionsController) viewWorktreeOptions(ref string) error {
+	return self.c.Helpers().Worktree.ViewWorktreeOptions(self.context, ref)
+}
diff --git a/pkg/gui/controllers/worktrees_controller.go b/pkg/gui/controllers/worktrees_controller.go
index 53dd2b32e..420d9ed39 100644
--- a/pkg/gui/controllers/worktrees_controller.go
+++ b/pkg/gui/controllers/worktrees_controller.go
@@ -62,7 +62,7 @@ func (self *WorktreesController) GetOnRenderToMain() func() error {
 			}
 
 			missing := ""
-			if self.c.Git().Worktree.IsWorktreePathMissing(worktree) {
+			if self.c.Git().Worktree.IsWorktreePathMissing(worktree.Path) {
 				missing = style.FgRed.Sprintf(" %s", self.c.Tr.MissingWorktree)
 			}
 
@@ -95,7 +95,7 @@ func (self *WorktreesController) remove(worktree *models.Worktree) error {
 		return self.c.ErrorMsg(self.c.Tr.CantDeleteMainWorktree)
 	}
 
-	if self.c.Git().Worktree.IsCurrentWorktree(worktree) {
+	if self.c.Git().Worktree.IsCurrentWorktree(worktree.Path) {
 		return self.c.ErrorMsg(self.c.Tr.CantDeleteCurrentWorktree)
 	}
 
@@ -107,7 +107,7 @@ func (self *WorktreesController) GetOnClick() func() error {
 }
 
 func (self *WorktreesController) enter(worktree *models.Worktree) error {
-	return self.c.Helpers().Worktree.Switch(worktree, context.WORKTREES_CONTEXT_KEY)
+	return self.c.Helpers().Worktree.Switch(worktree.Path, context.WORKTREES_CONTEXT_KEY)
 }
 
 func (self *WorktreesController) checkSelected(callback func(worktree *models.Worktree) error) func() error {
diff --git a/pkg/gui/presentation/worktrees.go b/pkg/gui/presentation/worktrees.go
index 1161afc44..4676a2847 100644
--- a/pkg/gui/presentation/worktrees.go
+++ b/pkg/gui/presentation/worktrees.go
@@ -8,11 +8,11 @@ import (
 	"github.com/samber/lo"
 )
 
-func GetWorktreeDisplayStrings(worktrees []*models.Worktree, isCurrent func(*models.Worktree) bool, isMissing func(*models.Worktree) bool) [][]string {
+func GetWorktreeDisplayStrings(worktrees []*models.Worktree, isCurrent func(string) bool, isMissing func(string) bool) [][]string {
 	return lo.Map(worktrees, func(worktree *models.Worktree, _ int) []string {
 		return GetWorktreeDisplayString(
-			isCurrent(worktree),
-			isMissing(worktree),
+			isCurrent(worktree.Path),
+			isMissing(worktree.Path),
 			worktree)
 	})
 }