From aa12ca6202d85f8eb28780100d7055940496d4cd Mon Sep 17 00:00:00 2001
From: Karol Musur <kamusur@gmail.com>
Date: Sun, 23 Oct 2022 19:35:33 +0200
Subject: [PATCH] feat: support git worktrees

`pro` now works in external git worktrees.
---
 repository/errors.go     |  5 ++-
 repository/repository.go | 81 ++++++++++++++++++++++++++++++++++------
 2 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/repository/errors.go b/repository/errors.go
index e2e578f..653d54b 100644
--- a/repository/errors.go
+++ b/repository/errors.go
@@ -3,6 +3,7 @@ package repository
 import "errors"
 
 var (
-	ErrNoActiveBranch = errors.New("no active branch")
-	ErrNoRemoteOrigin = errors.New("no remote named \"origin\" found")
+	ErrNoActiveBranch      = errors.New("no active branch")
+	ErrNoRemoteOrigin      = errors.New("no remote named \"origin\" found")
+	ErrUnableToReadGitFile = errors.New("unable to read .git file")
 )
diff --git a/repository/repository.go b/repository/repository.go
index 4674e94..3a906e5 100644
--- a/repository/repository.go
+++ b/repository/repository.go
@@ -2,13 +2,25 @@ package repository
 
 import (
 	"errors"
+	"os"
 	"path/filepath"
+	"strings"
 
+	"github.com/go-git/go-billy/v5/helper/chroot"
 	"github.com/go-git/go-git/v5"
+	"github.com/go-git/go-git/v5/storage/filesystem"
 )
 
 type Repository struct {
-	goGitRepository *git.Repository
+	workingDirectory string
+
+	// Root git directory, usually workingDirectory/.git
+	gitDirectory string
+
+	// Will be different from gitDirectory if the repository is an external worktree
+	// e.g. sample-repo/.git/worktrees/sample-repo-external
+	// See: https://git-scm.com/docs/git-worktree
+	worktreeGitDirectory string
 }
 
 // Return git repository in given directory or parent directories.
@@ -20,10 +32,10 @@ func FindInParents(path string) (Repository, error) {
 		return Repository{}, err
 	}
 
-	goGitRepository, err := git.PlainOpen(absolutePath)
-
+	_, err = git.PlainOpen(absolutePath)
+	// Found valid git repository
 	if err == nil {
-		return Repository{goGitRepository: goGitRepository}, nil
+		return makeRepository(absolutePath)
 	}
 
 	if errors.Is(err, git.ErrRepositoryNotExists) {
@@ -39,24 +51,71 @@ func FindInParents(path string) (Repository, error) {
 	return Repository{}, err
 }
 
+func makeRepository(workingDirectory string) (Repository, error) {
+	goGitRepository, err := git.PlainOpen(workingDirectory)
+	if err != nil {
+		return Repository{}, err
+	}
+
+	// Get git directory
+	storage, ok := goGitRepository.Storer.(*filesystem.Storage)
+	if !ok {
+		return Repository{}, errors.New("storage is not filesystem")
+	}
+	filesystem, ok := storage.Filesystem().(*chroot.ChrootHelper)
+	if !ok {
+		return Repository{}, errors.New("filesystem is not ChrootHelper")
+	}
+	gitDirectory := filesystem.Root()
+
+	// check if gitDirectory includes .git/worktrees in path
+	// if so, we are in an external worktree
+	if strings.Contains(gitDirectory, ".git/worktrees") {
+		worktreeGitDirectory := gitDirectory
+		// go 2 directories up to get the actual git directory
+		gitDirectory = filepath.Dir(filepath.Dir(gitDirectory))
+
+		return Repository{
+			workingDirectory:     workingDirectory,
+			gitDirectory:         gitDirectory,
+			worktreeGitDirectory: worktreeGitDirectory,
+		}, nil
+	} else {
+		return Repository{
+			workingDirectory:     workingDirectory,
+			gitDirectory:         gitDirectory,
+			worktreeGitDirectory: gitDirectory,
+		}, nil
+	}
+}
+
 func (repo *Repository) CurrentBranchName() (string, error) {
-	// get current head
-	head, err := repo.goGitRepository.Head()
+	// Get HEAD file contents from git directory
+	// We can't use go-git to get the current branch because it doesn't support worktrees
+	headFile, err := os.ReadFile(filepath.Join(repo.worktreeGitDirectory, "HEAD"))
 	if err != nil {
-		return "", err
+		return "", errors.New("unable to read HEAD")
 	}
 
-	if !head.Name().IsBranch() {
+	// Parse HEAD file to get branch name
+	headFileSplit := strings.Split(string(headFile), "ref: refs/heads/")
+	if len(headFileSplit) != 2 {
+		// if the HEAD file doesn't include "ref: refs/heads/", we are not on a branch
 		return "", ErrNoActiveBranch
 	}
+	branch := strings.TrimSpace(headFileSplit[1])
 
-	// current branch name
-	return head.Name().Short(), nil
+	return branch, nil
 }
 
 func (repo *Repository) OriginUrl() (string, error) {
+	goGitRepo, err := git.PlainOpen(repo.gitDirectory)
+	if err != nil {
+		return "", err
+	}
+
 	// check if there is a remote named origin
-	origin, err := repo.goGitRepository.Remote("origin")
+	origin, err := goGitRepo.Remote("origin")
 	if err != nil {
 		return "", ErrNoRemoteOrigin
 	}