Add ahead/behind information for @{push}

In a triangular workflow the branch that you're pulling from is not the same as
the one that you are pushing to. For example, some people find it useful to set
the upstream branch to origin/master so that pulling effectively rebases onto
master, and set the push.default git config to "current" so that "feature"
pushes to origin/feature.

Another example is a fork-based workflow where "feature" has upstream set to
upstream/main, and the repo has remote.pushDefault set to "origin", so pushing
on "feature" pushes to origin/feature.

This commit adds new fields to models.Branch that store the ahead/behind
information against the push branch; for the "normal" workflow where you pull
and push from/to the upstream branch, AheadForPush/BehindForPush will be the
same as AheadForPull/BehindForPull.

6 files changed, 48 insertions, 15 deletions

diff --git a/pkg/commands/git.go b/pkg/commands/git.go
index b43c8c4e5..7e7d9354f 100644
--- a/pkg/commands/git.go
+++ b/pkg/commands/git.go
@@ -134,7 +134,7 @@ func NewGitCommandAux(
 	worktreeCommands := git_commands.NewWorktreeCommands(gitCommon)
 	blameCommands := git_commands.NewBlameCommands(gitCommon)
 
-	branchLoader := git_commands.NewBranchLoader(cmn, cmd, branchCommands.CurrentBranchInfo, configCommands)
+	branchLoader := git_commands.NewBranchLoader(cmn, gitCommon, cmd, branchCommands.CurrentBranchInfo, configCommands)
 	commitFileLoader := git_commands.NewCommitFileLoader(cmn, cmd)
 	commitLoader := git_commands.NewCommitLoader(cmn, cmd, statusCommands.RebaseMode, gitCommon)
 	reflogCommitLoader := git_commands.NewReflogCommitLoader(cmn, cmd)
diff --git a/pkg/commands/git_commands/branch_loader.go b/pkg/commands/git_commands/branch_loader.go
index f1c68cd53..16777243a 100644
--- a/pkg/commands/git_commands/branch_loader.go
+++ b/pkg/commands/git_commands/branch_loader.go
@@ -40,6 +40,7 @@ type BranchInfo struct {
 // BranchLoader returns a list of Branch objects for the current repo
 type BranchLoader struct {
 	*common.Common
+	*GitCommon
 	cmd                  oscommands.ICmdObjBuilder
 	getCurrentBranchInfo func() (BranchInfo, error)
 	config               BranchLoaderConfigCommands
@@ -47,12 +48,14 @@ type BranchLoader struct {
 
 func NewBranchLoader(
 	cmn *common.Common,
+	gitCommon *GitCommon,
 	cmd oscommands.ICmdObjBuilder,
 	getCurrentBranchInfo func() (BranchInfo, error),
 	config BranchLoaderConfigCommands,
 ) *BranchLoader {
 	return &BranchLoader{
 		Common:               cmn,
+		GitCommon:            gitCommon,
 		cmd:                  cmd,
 		getCurrentBranchInfo: getCurrentBranchInfo,
 		config:               config,
@@ -61,7 +64,7 @@ func NewBranchLoader(
 
 // Load the list of branches for the current repo
 func (self *BranchLoader) Load(reflogCommits []*models.Commit) ([]*models.Branch, error) {
-	branches := self.obtainBranches()
+	branches := self.obtainBranches(self.version.IsAtLeast(2, 22, 0))
 
 	if self.AppState.LocalBranchSortOrder == "recency" {
 		reflogBranches := self.obtainReflogBranches(reflogCommits)
@@ -124,7 +127,7 @@ func (self *BranchLoader) Load(reflogCommits []*models.Commit) ([]*models.Branch
 	return branches, nil
 }
 
-func (self *BranchLoader) obtainBranches() []*models.Branch {
+func (self *BranchLoader) obtainBranches(canUsePushTrack bool) []*models.Branch {
 	output, err := self.getRawBranches()
 	if err != nil {
 		panic(err)
@@ -147,7 +150,7 @@ func (self *BranchLoader) obtainBranches() []*models.Branch {
 		}
 
 		storeCommitDateAsRecency := self.AppState.LocalBranchSortOrder != "recency"
-		return obtainBranch(split, storeCommitDateAsRecency), true
+		return obtainBranch(split, storeCommitDateAsRecency, canUsePushTrack), true
 	})
 }
 
@@ -183,23 +186,31 @@ var branchFields = []string{
 	"refname:short",
 	"upstream:short",
 	"upstream:track",
+	"push:track",
 	"subject",
 	"objectname",
 	"committerdate:unix",
 }
 
 // Obtain branch information from parsed line output of getRawBranches()
-func obtainBranch(split []string, storeCommitDateAsRecency bool) *models.Branch {
+func obtainBranch(split []string, storeCommitDateAsRecency bool, canUsePushTrack bool) *models.Branch {
 	headMarker := split[0]
 	fullName := split[1]
 	upstreamName := split[2]
 	track := split[3]
-	subject := split[4]
-	commitHash := split[5]
-	commitDate := split[6]
+	pushTrack := split[4]
+	subject := split[5]
+	commitHash := split[6]
+	commitDate := split[7]
 
 	name := strings.TrimPrefix(fullName, "heads/")
 	aheadForPull, behindForPull, gone := parseUpstreamInfo(upstreamName, track)
+	var aheadForPush, behindForPush string
+	if canUsePushTrack {
+		aheadForPush, behindForPush, _ = parseUpstreamInfo(upstreamName, pushTrack)
+	} else {
+		aheadForPush, behindForPush = aheadForPull, behindForPull
+	}
 
 	recency := ""
 	if storeCommitDateAsRecency {
@@ -213,6 +224,8 @@ func obtainBranch(split []string, storeCommitDateAsRecency bool) *models.Branch
 		Recency:       recency,
 		AheadForPull:  aheadForPull,
 		BehindForPull: behindForPull,
+		AheadForPush:  aheadForPush,
+		BehindForPush: behindForPush,
 		UpstreamGone:  gone,
 		Head:          headMarker == "*",
 		Subject:       subject,
diff --git a/pkg/commands/git_commands/branch_loader_test.go b/pkg/commands/git_commands/branch_loader_test.go
index e9b8001b5..2236374e5 100644
--- a/pkg/commands/git_commands/branch_loader_test.go
+++ b/pkg/commands/git_commands/branch_loader_test.go
@@ -25,12 +25,14 @@ func TestObtainBranch(t *testing.T) {
 	scenarios := []scenario{
 		{
 			testName:                 "TrimHeads",
-			input:                    []string{"", "heads/a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"", "heads/a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
 				Name:          "a_branch",
 				AheadForPull:  "?",
 				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
 				Head:          false,
 				Subject:       "subject",
 				CommitHash:    "123",
@@ -38,12 +40,14 @@ func TestObtainBranch(t *testing.T) {
 		},
 		{
 			testName:                 "NoUpstream",
-			input:                    []string{"", "a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
 				Name:          "a_branch",
 				AheadForPull:  "?",
 				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
 				Head:          false,
 				Subject:       "subject",
 				CommitHash:    "123",
@@ -51,12 +55,14 @@ func TestObtainBranch(t *testing.T) {
 		},
 		{
 			testName:                 "IsHead",
-			input:                    []string{"*", "a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"*", "a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
 				Name:          "a_branch",
 				AheadForPull:  "?",
 				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
 				Head:          true,
 				Subject:       "subject",
 				CommitHash:    "123",
@@ -64,12 +70,14 @@ func TestObtainBranch(t *testing.T) {
 		},
 		{
 			testName:                 "IsBehindAndAhead",
-			input:                    []string{"", "a_branch", "a_remote/a_branch", "[behind 2, ahead 3]", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "a_remote/a_branch", "[behind 2, ahead 3]", "[behind 2, ahead 3]", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
 				Name:          "a_branch",
 				AheadForPull:  "3",
 				BehindForPull: "2",
+				AheadForPush:  "3",
+				BehindForPush: "2",
 				Head:          false,
 				Subject:       "subject",
 				CommitHash:    "123",
@@ -77,13 +85,15 @@ func TestObtainBranch(t *testing.T) {
 		},
 		{
 			testName:                 "RemoteBranchIsGone",
-			input:                    []string{"", "a_branch", "a_remote/a_branch", "[gone]", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "a_remote/a_branch", "[gone]", "[gone]", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: false,
 			expectedBranch: &models.Branch{
 				Name:          "a_branch",
 				UpstreamGone:  true,
 				AheadForPull:  "?",
 				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
 				Head:          false,
 				Subject:       "subject",
 				CommitHash:    "123",
@@ -91,13 +101,15 @@ func TestObtainBranch(t *testing.T) {
 		},
 		{
 			testName:                 "WithCommitDateAsRecency",
-			input:                    []string{"", "a_branch", "", "", "subject", "123", timeStamp},
+			input:                    []string{"", "a_branch", "", "", "", "subject", "123", timeStamp},
 			storeCommitDateAsRecency: true,
 			expectedBranch: &models.Branch{
 				Name:          "a_branch",
 				Recency:       "2h",
 				AheadForPull:  "?",
 				BehindForPull: "?",
+				AheadForPush:  "?",
+				BehindForPush: "?",
 				Head:          false,
 				Subject:       "subject",
 				CommitHash:    "123",
@@ -107,7 +119,7 @@ func TestObtainBranch(t *testing.T) {
 
 	for _, s := range scenarios {
 		t.Run(s.testName, func(t *testing.T) {
-			branch := obtainBranch(s.input, s.storeCommitDateAsRecency)
+			branch := obtainBranch(s.input, s.storeCommitDateAsRecency, true)
 			assert.EqualValues(t, s.expectedBranch, branch)
 		})
 	}
diff --git a/pkg/commands/models/branch.go b/pkg/commands/models/branch.go
index 9e4772410..75d99f4bc 100644
--- a/pkg/commands/models/branch.go
+++ b/pkg/commands/models/branch.go
@@ -14,6 +14,10 @@ type Branch struct {
 	AheadForPull string
 	// how many commits behind we are from the remote branch (how many commits we can pull)
 	BehindForPull string
+	// how many commits ahead we are from the branch we're pushing to (which might not be the same as our upstream branch in a triangular workflow)
+	AheadForPush string
+	// how many commits behind we are from the branch we're pushing to (which might not be the same as our upstream branch in a triangular workflow)
+	BehindForPush string
 	// whether the remote branch is 'gone' i.e. we're tracking a remote branch that has been deleted
 	UpstreamGone bool
 	// whether this is the current branch. Exactly one branch should have this be true
diff --git a/pkg/gui/services/custom_commands/models.go b/pkg/gui/services/custom_commands/models.go
index bd3a69224..261bace45 100644
--- a/pkg/gui/services/custom_commands/models.go
+++ b/pkg/gui/services/custom_commands/models.go
@@ -51,6 +51,8 @@ type Branch struct {
 	Pullables      string // deprecated: use BehindForPull
 	AheadForPull   string
 	BehindForPull  string
+	AheadForPush   string
+	BehindForPush  string
 	UpstreamGone   bool
 	Head           bool
 	DetachedHead   bool
diff --git a/pkg/gui/services/custom_commands/session_state_loader.go b/pkg/gui/services/custom_commands/session_state_loader.go
index 81d30e5bf..6f39c5f8c 100644
--- a/pkg/gui/services/custom_commands/session_state_loader.go
+++ b/pkg/gui/services/custom_commands/session_state_loader.go
@@ -75,6 +75,8 @@ func branchShimFromModelBranch(branch *models.Branch) *Branch {
 		Pullables:      branch.BehindForPull,
 		AheadForPull:   branch.AheadForPull,
 		BehindForPull:  branch.BehindForPull,
+		AheadForPush:   branch.AheadForPush,
+		BehindForPush:  branch.BehindForPush,
 		UpstreamGone:   branch.UpstreamGone,
 		Head:           branch.Head,
 		DetachedHead:   branch.DetachedHead,