Associate branches with worktrees even when mid-rebase

3 files changed, 15 insertions, 15 deletions

diff --git a/pkg/gui/context/branches_context.go b/pkg/gui/context/branches_context.go
index 497b3a2c4..85c45c64a 100644
--- a/pkg/gui/context/branches_context.go
+++ b/pkg/gui/context/branches_context.go
@@ -31,6 +31,7 @@ func NewBranchesContext(c *ContextCommon) *BranchesContext {
 			c.Modes().Diffing.Ref,
 			c.Tr,
 			c.UserConfig,
+			c.Model().Worktrees,
 		)
 	}
 
diff --git a/pkg/gui/controllers/branches_controller.go b/pkg/gui/controllers/branches_controller.go
index 3ead83f35..4ea8099b2 100644
--- a/pkg/gui/controllers/branches_controller.go
+++ b/pkg/gui/controllers/branches_controller.go
@@ -202,11 +202,9 @@ func (self *BranchesController) press(selectedBranch *models.Branch) error {
 		return self.c.ErrorMsg(self.c.Tr.AlreadyCheckedOutBranch)
 	}
 
-	if selectedBranch.CheckedOutByOtherWorktree {
-		worktreeForRef, ok := self.worktreeForBranch(selectedBranch)
-		if ok && !self.c.Git().Worktree.IsCurrentWorktree(worktreeForRef.Path) {
-			return self.promptToCheckoutWorktree(worktreeForRef)
-		}
+	worktreeForRef, ok := self.worktreeForBranch(selectedBranch)
+	if ok && !self.c.Git().Worktree.IsCurrentWorktree(worktreeForRef.Path) {
+		return self.promptToCheckoutWorktree(worktreeForRef)
 	}
 
 	self.c.LogAction(self.c.Tr.Actions.CheckoutBranch)
@@ -214,13 +212,7 @@ func (self *BranchesController) press(selectedBranch *models.Branch) error {
 }
 
 func (self *BranchesController) worktreeForBranch(branch *models.Branch) (*models.Worktree, bool) {
-	for _, worktree := range self.c.Model().Worktrees {
-		if worktree.Branch == branch.Name {
-			return worktree, true
-		}
-	}
-
-	return nil, false
+	return git_commands.WorktreeForBranch(branch, self.c.Model().Worktrees)
 }
 
 func (self *BranchesController) promptToCheckoutWorktree(worktree *models.Worktree) error {
@@ -326,13 +318,17 @@ func (self *BranchesController) delete(branch *models.Branch) error {
 		return self.c.ErrorMsg(self.c.Tr.CantDeleteCheckOutBranch)
 	}
 
-	if branch.CheckedOutByOtherWorktree {
+	if self.checkedOutByOtherWorktree(branch) {
 		return self.promptWorktreeBranchDelete(branch)
 	}
 
 	return self.deleteWithForce(branch, false)
 }
 
+func (self *BranchesController) checkedOutByOtherWorktree(branch *models.Branch) bool {
+	return git_commands.CheckedOutByOtherWorktree(branch, self.c.Model().Worktrees)
+}
+
 func (self *BranchesController) promptWorktreeBranchDelete(selectedBranch *models.Branch) error {
 	worktree, ok := self.worktreeForBranch(selectedBranch)
 	if !ok {
diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go
index 42cbdf22b..6ad4e498e 100644
--- a/pkg/gui/presentation/branches.go
+++ b/pkg/gui/presentation/branches.go
@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/jesseduffield/generics/slices"
+	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/config"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation/icons"
@@ -23,10 +24,11 @@ func GetBranchListDisplayStrings(
 	diffName string,
 	tr *i18n.TranslationSet,
 	userConfig *config.UserConfig,
+	worktrees []*models.Worktree,
 ) [][]string {
 	return slices.Map(branches, func(branch *models.Branch) []string {
 		diffed := branch.Name == diffName
-		return getBranchDisplayStrings(branch, fullDescription, diffed, tr, userConfig)
+		return getBranchDisplayStrings(branch, fullDescription, diffed, tr, userConfig, worktrees)
 	})
 }
 
@@ -37,6 +39,7 @@ func getBranchDisplayStrings(
 	diffed bool,
 	tr *i18n.TranslationSet,
 	userConfig *config.UserConfig,
+	worktrees []*models.Worktree,
 ) []string {
 	displayName := b.Name
 	if b.DisplayName != "" {
@@ -50,7 +53,7 @@ func getBranchDisplayStrings(
 
 	coloredName := nameTextStyle.Sprint(displayName)
 	branchStatus := utils.WithPadding(ColoredBranchStatus(b, tr), 2, utils.AlignLeft)
-	if b.CheckedOutByOtherWorktree {
+	if git_commands.CheckedOutByOtherWorktree(b, worktrees) {
 		worktreeIcon := lo.Ternary(icons.IsIconEnabled(), icons.LINKED_WORKTREE_ICON, fmt.Sprintf("(%s)", tr.LcWorktree))
 		coloredName = fmt.Sprintf("%s %s", coloredName, style.FgDefault.Sprint(worktreeIcon))
 	}