Refactor

4 files changed, 25 insertions, 14 deletions

diff --git a/pkg/gui/context/worktrees_context.go b/pkg/gui/context/worktrees_context.go
index 044242225..ff3b833b5 100644
--- a/pkg/gui/context/worktrees_context.go
+++ b/pkg/gui/context/worktrees_context.go
@@ -4,7 +4,6 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
-	"github.com/samber/lo"
 )
 
 type WorktreesContext struct {
@@ -23,12 +22,11 @@ func NewWorktreesContext(c *ContextCommon) *WorktreesContext {
 	)
 
 	getDisplayStrings := func(startIdx int, length int) [][]string {
-		return lo.Map(c.Model().Worktrees, func(worktree *models.Worktree, _ int) []string {
-			return presentation.GetWorktreeDisplayString(
-				c.Git().Worktree.IsCurrentWorktree(worktree),
-				c.Git().Worktree.IsWorktreePathMissing(worktree),
-				worktree)
-		})
+		return presentation.GetWorktreeDisplayStrings(
+			c.Model().Worktrees,
+			c.Git().Worktree.IsCurrentWorktree,
+			c.Git().Worktree.IsWorktreePathMissing,
+		)
 	}
 
 	return &WorktreesContext{
diff --git a/pkg/gui/presentation/branches.go b/pkg/gui/presentation/branches.go
index cf0802d6f..243c9d0a3 100644
--- a/pkg/gui/presentation/branches.go
+++ b/pkg/gui/presentation/branches.go
@@ -12,6 +12,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/i18n"
 	"github.com/jesseduffield/lazygit/pkg/theme"
 	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/samber/lo"
 )
 
 var branchPrefixColorCache = make(map[string]style.TextStyle)
@@ -49,6 +50,10 @@ func getBranchDisplayStrings(
 
 	coloredName := nameTextStyle.Sprint(displayName)
 	branchStatus := utils.WithPadding(ColoredBranchStatus(b, tr), 2, utils.AlignLeft)
+	if b.CheckedOutByOtherWorktree {
+		worktreeIcon := lo.Ternary(icons.IsIconEnabled(), icons.LINKED_WORKTREE_ICON, "(worktree)")
+		coloredName = fmt.Sprintf("%s %s", coloredName, style.FgDefault.Sprint(worktreeIcon))
+	}
 	coloredName = fmt.Sprintf("%s %s", coloredName, branchStatus)
 
 	recencyColor := style.FgCyan
@@ -58,6 +63,7 @@ func getBranchDisplayStrings(
 
 	res := make([]string, 0, 6)
 	res = append(res, recencyColor.Sprint(b.Recency))
+
 	if icons.IsIconEnabled() {
 		res = append(res, nameTextStyle.Sprint(icons.IconForBranch(b)))
 	}
diff --git a/pkg/gui/presentation/icons/git_icons.go b/pkg/gui/presentation/icons/git_icons.go
index f6b56284e..17877c67e 100644
--- a/pkg/gui/presentation/icons/git_icons.go
+++ b/pkg/gui/presentation/icons/git_icons.go
@@ -71,10 +71,7 @@ func IconForStash(stash *models.StashEntry) string {
 	return STASH_ICON
 }
 
-func IconForWorktree(worktree *models.Worktree, missing bool) string {
-	if worktree.Main() {
-		return ""
-	}
+func IconForWorktree(missing bool) string {
 	if missing {
 		return MISSING_LINKED_WORKTREE_ICON
 	}
diff --git a/pkg/gui/presentation/worktrees.go b/pkg/gui/presentation/worktrees.go
index 063278ec6..55c1da08b 100644
--- a/pkg/gui/presentation/worktrees.go
+++ b/pkg/gui/presentation/worktrees.go
@@ -5,8 +5,18 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation/icons"
 	"github.com/jesseduffield/lazygit/pkg/gui/style"
 	"github.com/jesseduffield/lazygit/pkg/theme"
+	"github.com/samber/lo"
 )
 
+func GetWorktreeDisplayStrings(worktrees []*models.Worktree, isCurrent func(*models.Worktree) bool, isMissing func(*models.Worktree) bool) [][]string {
+	return lo.Map(worktrees, func(worktree *models.Worktree, _ int) []string {
+		return GetWorktreeDisplayString(
+			isCurrent(worktree),
+			isMissing(worktree),
+			worktree)
+	})
+}
+
 func GetWorktreeDisplayString(isCurrent bool, isPathMissing bool, worktree *models.Worktree) []string {
 	textStyle := theme.DefaultTextColor
 
@@ -17,13 +27,13 @@ func GetWorktreeDisplayString(isCurrent bool, isPathMissing bool, worktree *mode
 		currentColor = style.FgGreen
 	}
 
-	icon := icons.IconForWorktree(worktree, false)
+	icon := icons.IconForWorktree(false)
 	if isPathMissing {
 		textStyle = style.FgRed
-		icon = icons.IconForWorktree(worktree, true)
+		icon = icons.IconForWorktree(true)
 	}
 
-	res := make([]string, 0, 3)
+	res := []string{}
 	res = append(res, currentColor.Sprint(current))
 	if icons.IsIconEnabled() {
 		res = append(res, textStyle.Sprint(icon))