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