Associate branches with worktrees even when mid-rebase

pkg/commands/git_commands/branch_loader.go

@@ -2,7 +2,6 @@ package git_commands
 
 import (
 	"fmt"
-	"os"
 	"regexp"
 	"strings"
 
@@ -118,11 +117,6 @@ outer:
 }
 
 func (self *BranchLoader) obtainBranches() []*models.Branch {
-	currentDir, err := os.Getwd()
-	if err != nil {
-		panic(err)
-	}
-
 	output, err := self.getRawBranches()
 	if err != nil {
 		panic(err)
@@ -144,7 +138,7 @@ func (self *BranchLoader) obtainBranches() []*models.Branch {
 			return nil, false
 		}
 
-		return obtainBranch(split, currentDir), true
+		return obtainBranch(split), true
 	})
 }
 
@@ -172,32 +166,28 @@ var branchFields = []string{
 	"upstream:track",
 	"subject",
 	fmt.Sprintf("objectname:short=%d", utils.COMMIT_HASH_SHORT_SIZE),
-	"worktreepath",
 }
 
 // Obtain branch information from parsed line output of getRawBranches()
-func obtainBranch(split []string, currentDir string) *models.Branch {
+func obtainBranch(split []string) *models.Branch {
 	headMarker := split[0]
 	fullName := split[1]
 	upstreamName := split[2]
 	track := split[3]
 	subject := split[4]
 	commitHash := split[5]
-	branchDir := split[6]
-	checkedOutByOtherWorktree := len(branchDir) > 0 && branchDir != currentDir
 
 	name := strings.TrimPrefix(fullName, "heads/")
 	pushables, pullables, gone := parseUpstreamInfo(upstreamName, track)
 
 	return &models.Branch{
-		Name:                      name,
-		Pushables:                 pushables,
-		Pullables:                 pullables,
-		UpstreamGone:              gone,
-		Head:                      headMarker == "*",
-		Subject:                   subject,
-		CommitHash:                commitHash,
-		CheckedOutByOtherWorktree: checkedOutByOtherWorktree,
+		Name:         name,
+		Pushables:    pushables,
+		Pullables:    pullables,
+		UpstreamGone: gone,
+		Head:         headMarker == "*",
+		Subject:      subject,
+		CommitHash:   commitHash,
 	}
 }
 

pkg/commands/git_commands/branch_loader_test.go

@@ -81,7 +81,7 @@ func TestObtainBranch(t *testing.T) {
 
 	for _, s := range scenarios {
 		t.Run(s.testName, func(t *testing.T) {
-			branch := obtainBranch(s.input, "current-dir")
+			branch := obtainBranch(s.input)
 			assert.EqualValues(t, s.expectedBranch, branch)
 		})
 	}

pkg/commands/git_commands/worktree.go

@@ -6,6 +6,8 @@ import (
 	"io/fs"
 	"log"
 	"os"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
 )
 
 type WorktreeCommands struct {
@@ -84,3 +86,22 @@ func (self *WorktreeCommands) IsWorktreePathMissing(path string) bool {
 func EqualPath(a string, b string) bool {
 	return a == b
 }
+
+func WorktreeForBranch(branch *models.Branch, worktrees []*models.Worktree) (*models.Worktree, bool) {
+	for _, worktree := range worktrees {
+		if worktree.Branch == branch.Name {
+			return worktree, true
+		}
+	}
+
+	return nil, false
+}
+
+func CheckedOutByOtherWorktree(branch *models.Branch, worktrees []*models.Worktree) bool {
+	worktree, ok := WorktreeForBranch(branch, worktrees)
+	if !ok {
+		return false
+	}
+
+	return !IsCurrentWorktree(worktree.Path)
+}

pkg/commands/git_commands/worktree_loader.go

@@ -2,6 +2,7 @@ package git_commands
 
 import (
 	"os"
+	"path/filepath"
 	"strings"
 
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
@@ -76,9 +77,66 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
 		}
 	}
 
+	// Some worktrees are on a branch but are mid-rebase, and in those cases,
+	// `git worktree list` will not show the branch name. We can get the branch
+	// name from the `rebase-merge/head-name` file (if it exists) in the folder
+	// for the worktree in the parent repo's .git/worktrees folder.
+	for _, worktree := range worktrees {
+		// No point checking if we already have a branch name
+		if worktree.Branch != "" {
+			continue
+		}
+
+		rebaseBranch, ok := rebaseBranch(worktree.Path)
+		if ok {
+			worktree.Branch = rebaseBranch
+		}
+	}
+
 	return worktrees, nil
 }
 
+func rebaseBranch(worktreePath string) (string, bool) {
+	// need to find the actual path of the worktree in the .git dir
+	gitPath, ok := worktreeGitPath(worktreePath)
+	if !ok {
+		return "", false
+	}
+
+	// now we look inside that git path for a file `rebase-merge/head-name`
+	// if it exists, we update the worktree to say that it has that for a head
+	headNameContents, err := os.ReadFile(filepath.Join(gitPath, "rebase-merge", "head-name"))
+	if err != nil {
+		return "", false
+	}
+
+	headName := strings.TrimSpace(string(headNameContents))
+	shortHeadName := strings.TrimPrefix(headName, "refs/heads/")
+
+	return shortHeadName, true
+}
+
+func worktreeGitPath(worktreePath string) (string, bool) {
+	// first we get the path of the worktree, then we look at the contents of the `.git` file in that path
+	// then we look for the line that says `gitdir: /path/to/.git/worktrees/<worktree-name>`
+	// then we return that path
+	gitFileContents, err := os.ReadFile(filepath.Join(worktreePath, ".git"))
+	if err != nil {
+		return "", false
+	}
+
+	gitDirLine := lo.Filter(strings.Split(string(gitFileContents), "\n"), func(line string, _ int) bool {
+		return strings.HasPrefix(line, "gitdir: ")
+	})
+
+	if len(gitDirLine) == 0 {
+		return "", false
+	}
+
+	gitDir := strings.TrimPrefix(gitDirLine[0], "gitdir: ")
+	return gitDir, true
+}
+
 type pathWithIndexT struct {
 	path  string
 	index int

pkg/commands/models/branch.go

@@ -26,8 +26,6 @@ type Branch struct {
 	Subject string
 	// commit hash
 	CommitHash string
-
-	CheckedOutByOtherWorktree bool
 }
 
 func (b *Branch) FullRefName() string {

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,
 		)
 	}
 

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 {

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))
 	}