Visualize local branch heads in commits panel

We want to mark all local branch heads with a "*" in the local commits panel, to
make it easier to see how branches are stacked onto each other. In order to not
confuse users with "*" markers that they don't understand, do this only for the
case where users actually use stacked branches; those users are likely not going
to be confused by the display. This means we want to filter out a few branch
heads that shouldn't get the marker: the current branch, any main branch, and
any old branch that has been merged to master already.

6 files changed, 122 insertions, 5 deletions

diff --git a/pkg/gui/context/local_commits_context.go b/pkg/gui/context/local_commits_context.go
index 84204591c..9da76104c 100644
--- a/pkg/gui/context/local_commits_context.go
+++ b/pkg/gui/context/local_commits_context.go
@@ -42,6 +42,8 @@ func NewLocalCommitsContext(c *ContextCommon) *LocalCommitsContext {
 		return presentation.GetCommitListDisplayStrings(
 			c.Common,
 			c.Model().Commits,
+			c.Model().Branches,
+			c.Model().CheckedOutBranch,
 			c.State().GetRepoState().GetScreenMode() != types.SCREEN_NORMAL,
 			c.Modes().CherryPicking.SelectedShaSet(),
 			c.Modes().Diffing.Ref,
diff --git a/pkg/gui/context/sub_commits_context.go b/pkg/gui/context/sub_commits_context.go
index 0cf884589..1a606ac91 100644
--- a/pkg/gui/context/sub_commits_context.go
+++ b/pkg/gui/context/sub_commits_context.go
@@ -37,6 +37,13 @@ func NewSubCommitsContext(
 	}
 
 	getDisplayStrings := func(startIdx int, length int) [][]string {
+		// This can happen if a sub-commits view is asked to be rerendered while
+		// it is invisble; for example when switching screen modes, which
+		// rerenders all views.
+		if viewModel.GetRef() == nil {
+			return [][]string{}
+		}
+
 		selectedCommitSha := ""
 		if c.CurrentContext().GetKey() == SUB_COMMITS_CONTEXT_KEY {
 			selectedCommit := viewModel.GetSelected()
@@ -47,6 +54,8 @@ func NewSubCommitsContext(
 		return presentation.GetCommitListDisplayStrings(
 			c.Common,
 			c.Model().SubCommits,
+			c.Model().Branches,
+			viewModel.GetRef().RefName(),
 			c.State().GetRepoState().GetScreenMode() != types.SCREEN_NORMAL,
 			c.Modes().CherryPicking.SelectedShaSet(),
 			c.Modes().Diffing.Ref,
diff --git a/pkg/gui/controllers/helpers/refresh_helper.go b/pkg/gui/controllers/helpers/refresh_helper.go
index b6da18374..fd586c4d5 100644
--- a/pkg/gui/controllers/helpers/refresh_helper.go
+++ b/pkg/gui/controllers/helpers/refresh_helper.go
@@ -439,6 +439,12 @@ func (self *RefreshHelper) refreshBranches() {
 		self.c.Log.Error(err)
 	}
 
+	// Need to re-render the commits view because the visualization of local
+	// branch heads might have changed
+	if err := self.c.Contexts().LocalCommits.HandleRender(); err != nil {
+		self.c.Log.Error(err)
+	}
+
 	self.refreshStatus()
 }
 
diff --git a/pkg/gui/presentation/commits.go b/pkg/gui/presentation/commits.go
index 53d3aaf0d..4affe6a57 100644
--- a/pkg/gui/presentation/commits.go
+++ b/pkg/gui/presentation/commits.go
@@ -39,6 +39,8 @@ type bisectBounds struct {
 func GetCommitListDisplayStrings(
 	common *common.Common,
 	commits []*models.Commit,
+	branches []*models.Branch,
+	currentBranchName string,
 	fullDescription bool,
 	cherryPickedCommitShaSet *set.Set[string],
 	diffName string,
@@ -99,6 +101,24 @@ func GetCommitListDisplayStrings(
 		getGraphLine = func(idx int) string { return "" }
 	}
 
+	// Determine the hashes of the local branches for which we want to show a
+	// branch marker in the commits list. We only want to do this for branches
+	// that are not the current branch, and not any of the main branches. The
+	// goal is to visualize stacks of local branches, so anything that doesn't
+	// contribute to a branch stack shouldn't show a marker.
+	branchHeadsToVisualize := set.NewFromSlice(lo.FilterMap(branches,
+		func(b *models.Branch, index int) (string, bool) {
+			return b.CommitHash,
+				// Don't consider branches that don't have a commit hash. As far
+				// as I can see, this happens for a detached head, so filter
+				// these out
+				b.CommitHash != "" &&
+					// Don't show a marker for the current branch
+					b.Name != currentBranchName &&
+					// Don't show a marker for main branches
+					!lo.Contains(common.UserConfig.Git.MainBranches, b.Name)
+		}))
+
 	lines := make([][]string, 0, len(filteredCommits))
 	var bisectStatus BisectStatus
 	for i, commit := range filteredCommits {
@@ -112,6 +132,7 @@ func GetCommitListDisplayStrings(
 		lines = append(lines, displayCommit(
 			common,
 			commit,
+			branchHeadsToVisualize,
 			cherryPickedCommitShaSet,
 			diffName,
 			timeFormat,
@@ -260,6 +281,7 @@ func getBisectStatusText(bisectStatus BisectStatus, bisectInfo *git_commands.Bis
 func displayCommit(
 	common *common.Common,
 	commit *models.Commit,
+	branchHeadsToVisualize *set.Set[string],
 	cherryPickedCommitShaSet *set.Set[string],
 	diffName string,
 	timeFormat string,
@@ -290,6 +312,11 @@ func displayCommit(
 		if len(commit.Tags) > 0 {
 			tagString = theme.DiffTerminalColor.SetBold().Sprint(strings.Join(commit.Tags, " ")) + " "
 		}
+
+		if branchHeadsToVisualize.Includes(commit.Sha) && commit.Status != models.StatusMerged {
+			tagString = style.FgCyan.SetBold().Sprint(
+				lo.Ternary(icons.IsIconEnabled(), icons.BRANCH_ICON, "*") + " " + tagString)
+		}
 	}
 
 	name := commit.Name
diff --git a/pkg/gui/presentation/commits_test.go b/pkg/gui/presentation/commits_test.go
index 31c6c813d..b6684f80e 100644
--- a/pkg/gui/presentation/commits_test.go
+++ b/pkg/gui/presentation/commits_test.go
@@ -28,6 +28,8 @@ func TestGetCommitListDisplayStrings(t *testing.T) {
 	scenarios := []struct {
 		testName                 string
 		commits                  []*models.Commit
+		branches                 []*models.Branch
+		currentBranchName        string
 		fullDescription          bool
 		cherryPickedCommitShaSet *set.Set[string]
 		diffName                 string
@@ -73,6 +75,73 @@ func TestGetCommitListDisplayStrings(t *testing.T) {
 						`),
 		},
 		{
+			testName: "commit with tags",
+			commits: []*models.Commit{
+				{Name: "commit1", Sha: "sha1", Tags: []string{"tag1", "tag2"}},
+				{Name: "commit2", Sha: "sha2"},
+			},
+			startIdx:                 0,
+			length:                   2,
+			showGraph:                false,
+			bisectInfo:               git_commands.NewNullBisectInfo(),
+			cherryPickedCommitShaSet: set.New[string](),
+			now:                      time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC),
+			expected: formatExpected(`
+		sha1 tag1 tag2 commit1
+		sha2 commit2
+						`),
+		},
+		{
+			testName: "show local branch head, except the current branch, main branches, or merged branches",
+			commits: []*models.Commit{
+				{Name: "commit1", Sha: "sha1"},
+				{Name: "commit2", Sha: "sha2"},
+				{Name: "commit3", Sha: "sha3"},
+				{Name: "commit4", Sha: "sha4", Status: models.StatusMerged},
+			},
+			branches: []*models.Branch{
+				{Name: "current-branch", CommitHash: "sha1", Head: true},
+				{Name: "other-branch", CommitHash: "sha2", Head: false},
+				{Name: "master", CommitHash: "sha3", Head: false},
+				{Name: "old-branch", CommitHash: "sha4", Head: false},
+			},
+			currentBranchName:        "current-branch",
+			startIdx:                 0,
+			length:                   4,
+			showGraph:                false,
+			bisectInfo:               git_commands.NewNullBisectInfo(),
+			cherryPickedCommitShaSet: set.New[string](),
+			now:                      time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC),
+			expected: formatExpected(`
+		sha1 commit1
+		sha2 * commit2
+		sha3 commit3
+		sha4 commit4
+						`),
+		},
+		{
+			testName: "show local branch head and tag if both exist",
+			commits: []*models.Commit{
+				{Name: "commit1", Sha: "sha1"},
+				{Name: "commit2", Sha: "sha2", Tags: []string{"some-tag"}},
+				{Name: "commit3", Sha: "sha3"},
+			},
+			branches: []*models.Branch{
+				{Name: "some-branch", CommitHash: "sha2"},
+			},
+			startIdx:                 0,
+			length:                   3,
+			showGraph:                false,
+			bisectInfo:               git_commands.NewNullBisectInfo(),
+			cherryPickedCommitShaSet: set.New[string](),
+			now:                      time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC),
+			expected: formatExpected(`
+		sha1 commit1
+		sha2 * some-tag commit2
+		sha3 commit3
+						`),
+		},
+		{
 			testName: "showing graph",
 			commits: []*models.Commit{
 				{Name: "commit1", Sha: "sha1", Parents: []string{"sha2", "sha3"}},
@@ -285,6 +354,8 @@ func TestGetCommitListDisplayStrings(t *testing.T) {
 				result := GetCommitListDisplayStrings(
 					common,
 					s.commits,
+					s.branches,
+					s.currentBranchName,
 					s.fullDescription,
 					s.cherryPickedCommitShaSet,
 					s.diffName,
diff --git a/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go b/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go
index 2302a3b69..3fa221d72 100644
--- a/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go
+++ b/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go
@@ -10,7 +10,9 @@ var DropTodoCommitWithUpdateRef = NewIntegrationTest(NewIntegrationTestArgs{
 	ExtraCmdArgs: []string{},
 	Skip:         false,
 	GitVersion:   AtLeast("2.38.0"),
-	SetupConfig:  func(config *config.AppConfig) {},
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Git.MainBranches = []string{"master"}
+	},
 	SetupRepo: func(shell *Shell) {
 		shell.
 			CreateNCommits(1).
@@ -28,7 +30,7 @@ var DropTodoCommitWithUpdateRef = NewIntegrationTest(NewIntegrationTestArgs{
 				Contains("CI commit 07").IsSelected(),
 				Contains("CI commit 06"),
 				Contains("CI commit 05"),
-				Contains("CI commit 04"),
+				Contains("CI * commit 04"),
 				Contains("CI commit 03"),
 				Contains("CI commit 02"),
 				Contains("CI commit 01"),
@@ -40,8 +42,8 @@ var DropTodoCommitWithUpdateRef = NewIntegrationTest(NewIntegrationTestArgs{
 				Contains("pick").Contains("CI commit 07"),
 				Contains("pick").Contains("CI commit 06"),
 				Contains("pick").Contains("CI commit 05"),
-				Contains("update-ref").Contains("branch1"),
-				Contains("pick").Contains("CI commit 04"),
+				Contains("update-ref").Contains("branch1").DoesNotContain("*"),
+				Contains("pick").Contains("CI * commit 04"),
 				Contains("pick").Contains("CI commit 03"),
 				Contains("<-- YOU ARE HERE --- commit 02"),
 				Contains("CI commit 01"),
@@ -56,7 +58,7 @@ var DropTodoCommitWithUpdateRef = NewIntegrationTest(NewIntegrationTestArgs{
 			Lines(
 				Contains("CI commit 07"),
 				Contains("CI commit 05"),
-				Contains("CI commit 04"),
+				Contains("CI * commit 04"),
 				Contains("CI commit 03"),
 				Contains("CI commit 02"),
 				Contains("CI commit 01"),