Factor out CommitLoader.mainBranches into its own class, and store it in Model

7 files changed, 116 insertions, 68 deletions

diff --git a/pkg/commands/git_commands/commit_loader.go b/pkg/commands/git_commands/commit_loader.go
index 86418453d..643a7c184 100644
--- a/pkg/commands/git_commands/commit_loader.go
+++ b/pkg/commands/git_commands/commit_loader.go
@@ -35,11 +35,6 @@ type CommitLoader struct {
 	readFile      func(filename string) ([]byte, error)
 	walkFiles     func(root string, fn filepath.WalkFunc) error
 	dotGitDir     string
-	// List of main branches that exist in the repo.
-	// We use these to obtain the merge base of the branch.
-	// When nil, we're yet to obtain the list of existing main branches.
-	// When an empty slice, we've obtained the list and it's empty.
-	mainBranches []string
 	*GitCommon
 }
 
@@ -56,7 +51,6 @@ func NewCommitLoader(
 		getRebaseMode: getRebaseMode,
 		readFile:      os.ReadFile,
 		walkFiles:     filepath.Walk,
-		mainBranches:  nil,
 		GitCommon:     gitCommon,
 	}
 }
@@ -72,6 +66,7 @@ type GetCommitsOptions struct {
 	All bool
 	// If non-empty, show divergence from this ref (left-right log)
 	RefToShowDivergenceFrom string
+	MainBranches            *MainBranches
 }
 
 // GetCommits obtains the commits of the current branch
@@ -108,9 +103,9 @@ func (self *CommitLoader) GetCommits(opts GetCommitsOptions) ([]*models.Commit,
 	go utils.Safe(func() {
 		defer wg.Done()
 
-		ancestor = self.getMergeBase(opts.RefName)
+		ancestor = self.getMergeBase(opts.RefName, opts.MainBranches)
 		if opts.RefToShowDivergenceFrom != "" {
-			remoteAncestor = self.getMergeBase(opts.RefToShowDivergenceFrom)
+			remoteAncestor = self.getMergeBase(opts.RefToShowDivergenceFrom, opts.MainBranches)
 		}
 	})
 
@@ -471,12 +466,9 @@ func setCommitMergedStatuses(ancestor string, commits []*models.Commit) {
 	}
 }
 
-func (self *CommitLoader) getMergeBase(refName string) string {
-	if self.mainBranches == nil {
-		self.mainBranches = self.getExistingMainBranches()
-	}
-
-	if len(self.mainBranches) == 0 {
+func (self *CommitLoader) getMergeBase(refName string, existingMainBranches *ExistingMainBranches) string {
+	mainBranches := existingMainBranches.Get()
+	if len(mainBranches) == 0 {
 		return ""
 	}
 
@@ -491,63 +483,12 @@ func (self *CommitLoader) getMergeBase(refName string) string {
 	// also not very common, but can totally happen and is not an error.
 
 	output, _ := self.cmd.New(
-		NewGitCmd("merge-base").Arg(refName).Arg(self.mainBranches...).
+		NewGitCmd("merge-base").Arg(refName).Arg(mainBranches...).
 			ToArgv(),
 	).DontLog().RunWithOutput()
 	return ignoringWarnings(output)
 }
 
-func (self *CommitLoader) getExistingMainBranches() []string {
-	var existingBranches []string
-	var wg sync.WaitGroup
-
-	mainBranches := self.UserConfig.Git.MainBranches
-	existingBranches = make([]string, len(mainBranches))
-
-	for i, branchName := range mainBranches {
-		wg.Add(1)
-		go utils.Safe(func() {
-			defer wg.Done()
-
-			// Try to determine upstream of local main branch
-			if ref, err := self.cmd.New(
-				NewGitCmd("rev-parse").Arg("--symbolic-full-name", branchName+"@{u}").ToArgv(),
-			).DontLog().RunWithOutput(); err == nil {
-				existingBranches[i] = strings.TrimSpace(ref)
-				return
-			}
-
-			// If this failed, a local branch for this main branch doesn't exist or it
-			// has no upstream configured. Try looking for one in the "origin" remote.
-			ref := "refs/remotes/origin/" + branchName
-			if err := self.cmd.New(
-				NewGitCmd("rev-parse").Arg("--verify", "--quiet", ref).ToArgv(),
-			).DontLog().Run(); err == nil {
-				existingBranches[i] = ref
-				return
-			}
-
-			// If this failed as well, try if we have the main branch as a local
-			// branch. This covers the case where somebody is using git locally
-			// for something, but never pushing anywhere.
-			ref = "refs/heads/" + branchName
-			if err := self.cmd.New(
-				NewGitCmd("rev-parse").Arg("--verify", "--quiet", ref).ToArgv(),
-			).DontLog().Run(); err == nil {
-				existingBranches[i] = ref
-			}
-		})
-	}
-
-	wg.Wait()
-
-	existingBranches = lo.Filter(existingBranches, func(branch string, _ int) bool {
-		return branch != ""
-	})
-
-	return existingBranches
-}
-
 func ignoringWarnings(commandOutput string) string {
 	trimmedOutput := strings.TrimSpace(commandOutput)
 	split := strings.Split(trimmedOutput, "\n")
diff --git a/pkg/commands/git_commands/commit_loader_test.go b/pkg/commands/git_commands/commit_loader_test.go
index fe4f39585..a8ef9e69a 100644
--- a/pkg/commands/git_commands/commit_loader_test.go
+++ b/pkg/commands/git_commands/commit_loader_test.go
@@ -307,10 +307,11 @@ func TestGetCommits(t *testing.T) {
 			common := utils.NewDummyCommon()
 			common.AppState = &config.AppState{}
 			common.AppState.GitLogOrder = scenario.logOrder
+			cmd := oscommands.NewDummyCmdObjBuilder(scenario.runner)
 
 			builder := &CommitLoader{
 				Common:        common,
-				cmd:           oscommands.NewDummyCmdObjBuilder(scenario.runner),
+				cmd:           cmd,
 				getRebaseMode: func() (enums.RebaseMode, error) { return scenario.rebaseMode, nil },
 				dotGitDir:     ".git",
 				readFile: func(filename string) ([]byte, error) {
@@ -322,7 +323,9 @@ func TestGetCommits(t *testing.T) {
 			}
 
 			common.UserConfig.Git.MainBranches = scenario.mainBranches
-			commits, err := builder.GetCommits(scenario.opts)
+			opts := scenario.opts
+			opts.MainBranches = NewMainBranches(scenario.mainBranches, cmd)
+			commits, err := builder.GetCommits(opts)
 
 			assert.Equal(t, scenario.expectedCommits, commits)
 			assert.Equal(t, scenario.expectedError, err)
diff --git a/pkg/commands/git_commands/main_branches.go b/pkg/commands/git_commands/main_branches.go
new file mode 100644
index 000000000..f8b8b85f2
--- /dev/null
+++ b/pkg/commands/git_commands/main_branches.go
@@ -0,0 +1,98 @@
+package git_commands
+
+import (
+	"strings"
+	"sync"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/oscommands"
+	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/samber/lo"
+	"github.com/sasha-s/go-deadlock"
+)
+
+type MainBranches struct {
+	// List of main branches configured by the user. Just the bare names.
+	configuredMainBranches []string
+	// Which of these actually exist in the repository. Full ref names, and it
+	// could be either "refs/heads/..." or "refs/remotes/origin/..." depending
+	// on which one exists for a given bare name.
+	existingMainBranches []string
+
+	cmd   oscommands.ICmdObjBuilder
+	mutex *deadlock.Mutex
+}
+
+func NewMainBranches(
+	configuredMainBranches []string,
+	cmd oscommands.ICmdObjBuilder,
+) *MainBranches {
+	return &MainBranches{
+		configuredMainBranches: configuredMainBranches,
+		existingMainBranches:   nil,
+		cmd:                    cmd,
+		mutex:                  &deadlock.Mutex{},
+	}
+}
+
+// Get the list of main branches that exist in the repository. This is a list of
+// full ref names.
+func (self *MainBranches) Get() []string {
+	self.mutex.Lock()
+	defer self.mutex.Unlock()
+
+	if self.existingMainBranches == nil {
+		self.existingMainBranches = self.determineMainBranches()
+	}
+
+	return self.existingMainBranches
+}
+
+func (self *MainBranches) determineMainBranches() []string {
+	var existingBranches []string
+	var wg sync.WaitGroup
+
+	existingBranches = make([]string, len(self.configuredMainBranches))
+
+	for i, branchName := range self.configuredMainBranches {
+		wg.Add(1)
+		go utils.Safe(func() {
+			defer wg.Done()
+
+			// Try to determine upstream of local main branch
+			if ref, err := self.cmd.New(
+				NewGitCmd("rev-parse").Arg("--symbolic-full-name", branchName+"@{u}").ToArgv(),
+			).DontLog().RunWithOutput(); err == nil {
+				existingBranches[i] = strings.TrimSpace(ref)
+				return
+			}
+
+			// If this failed, a local branch for this main branch doesn't exist or it
+			// has no upstream configured. Try looking for one in the "origin" remote.
+			ref := "refs/remotes/origin/" + branchName
+			if err := self.cmd.New(
+				NewGitCmd("rev-parse").Arg("--verify", "--quiet", ref).ToArgv(),
+			).DontLog().Run(); err == nil {
+				existingBranches[i] = ref
+				return
+			}
+
+			// If this failed as well, try if we have the main branch as a local
+			// branch. This covers the case where somebody is using git locally
+			// for something, but never pushing anywhere.
+			ref = "refs/heads/" + branchName
+			if err := self.cmd.New(
+				NewGitCmd("rev-parse").Arg("--verify", "--quiet", ref).ToArgv(),
+			).DontLog().Run(); err == nil {
+				existingBranches[i] = ref
+			}
+		})
+	}
+
+	wg.Wait()
+
+	existingBranches = lo.Filter(existingBranches, func(branch string, _ int) bool {
+		return branch != ""
+	})
+
+	return existingBranches
+}
diff --git a/pkg/gui/controllers/helpers/refresh_helper.go b/pkg/gui/controllers/helpers/refresh_helper.go
index b927296fc..02a26ded9 100644
--- a/pkg/gui/controllers/helpers/refresh_helper.go
+++ b/pkg/gui/controllers/helpers/refresh_helper.go
@@ -331,6 +331,7 @@ func (self *RefreshHelper) refreshCommitsWithLimit() error {
 			RefName:              self.refForLog(),
 			RefForPushedStatus:   checkedOutBranchName,
 			All:                  self.c.Contexts().LocalCommits.GetShowWholeGitGraph(),
+			MainBranches:         self.c.Model().MainBranches,
 		},
 	)
 	if err != nil {
@@ -357,6 +358,7 @@ func (self *RefreshHelper) refreshSubCommitsWithLimit() error {
 			RefName:                 self.c.Contexts().SubCommits.GetRef().FullRefName(),
 			RefToShowDivergenceFrom: self.c.Contexts().SubCommits.GetRefToShowDivergenceFrom(),
 			RefForPushedStatus:      self.c.Contexts().SubCommits.GetRef().FullRefName(),
+			MainBranches:            self.c.Model().MainBranches,
 		},
 	)
 	if err != nil {
diff --git a/pkg/gui/controllers/helpers/sub_commits_helper.go b/pkg/gui/controllers/helpers/sub_commits_helper.go
index c31d50937..f1cecf7f5 100644
--- a/pkg/gui/controllers/helpers/sub_commits_helper.go
+++ b/pkg/gui/controllers/helpers/sub_commits_helper.go
@@ -44,6 +44,7 @@ func (self *SubCommitsHelper) ViewSubCommits(opts ViewSubCommitsOpts) error {
 			RefName:                 opts.Ref.FullRefName(),
 			RefForPushedStatus:      opts.Ref.FullRefName(),
 			RefToShowDivergenceFrom: opts.RefToShowDivergenceFrom,
+			MainBranches:            self.c.Model().MainBranches,
 		},
 	)
 	if err != nil {
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index 0c0f36370..5f2fd55bf 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -379,6 +379,7 @@ func (gui *Gui) resetState(startArgs appTypes.StartArgs) types.Context {
 			BisectInfo:            git_commands.NewNullBisectInfo(),
 			FilesTrie:             patricia.NewTrie(),
 			Authors:               map[string]*models.Author{},
+			MainBranches:          git_commands.NewMainBranches(gui.UserConfig.Git.MainBranches, gui.os.Cmd),
 		},
 		Modes: &types.Modes{
 			Filtering:        filtering.New(startArgs.FilterPath, ""),
diff --git a/pkg/gui/types/common.go b/pkg/gui/types/common.go
index d50173078..77f2f56eb 100644
--- a/pkg/gui/types/common.go
+++ b/pkg/gui/types/common.go
@@ -281,6 +281,8 @@ type Model struct {
 	// we're on a detached head because we're rebasing or bisecting.
 	CheckedOutBranch string
 
+	MainBranches *git_commands.MainBranches
+
 	// for displaying suggestions while typing in a file name
 	FilesTrie *patricia.Trie