Count commits ahead and behind the upstream branch

git/repo_test.go

@@ -126,17 +126,30 @@ func newRepoWithBranchBehind(t *testing.T) *Repo {
 	return r
 }
 
+// generate repo with 2 commits ahead and 3 behind the origin
 func newRepoWithBranchAheadAndBehind(t *testing.T) *Repo {
 	origin := newRepoWithCommit(t)
 
 	r := origin.clone(t)
-	r.writeFile(t, "local.new", "I'm a new file on local")
+	r.writeFile(t, "local.new", "local 1")
 	r.addFile(t, "local.new")
-	r.newCommit(t, "new local commit")
+	r.newCommit(t, "1st local commit")
 
-	origin.writeFile(t, "origin.new", "I'm a new file on origin")
+	r.writeFile(t, "local.new", "local 2")
+	r.addFile(t, "local.new")
+	r.newCommit(t, "2nd local commit")
+
+	origin.writeFile(t, "origin.new", "origin 1")
 	origin.addFile(t, "origin.new")
-	origin.newCommit(t, "new origin commit")
+	origin.newCommit(t, "1st origin commit")
+
+	origin.writeFile(t, "origin.new", "origin 2")
+	origin.addFile(t, "origin.new")
+	origin.newCommit(t, "2nd origin commit")
+
+	origin.writeFile(t, "origin.new", "origin 3")
+	origin.addFile(t, "origin.new")
+	origin.newCommit(t, "3rd origin commit")
 
 	r.fetch(t)
 	return r

git/status.go

@@ -5,6 +5,8 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/go-git/go-git/v5/plumbing/revlist"
+
 	"github.com/spf13/viper"
 
 	"github.com/go-git/go-billy/v5/osfs"
@@ -35,10 +37,10 @@ type RepoStatus struct {
 }
 
 type BranchStatus struct {
-	Name      string
+	Name     string
-	Upstream  string
+	Upstream string
-	NeedsPull bool
+	Ahead    int
-	NeedsPush bool
+	Behind   int
 }
 
 func (r *Repo) LoadStatus() error {
@@ -172,14 +174,14 @@ func (r *Repo) newBranchStatus(branch string) (*BranchStatus, error) {
 		return bs, nil
 	}
 
-	needsPull, needsPush, err := r.needsPullOrPush(branch, upstream)
+	ahead, behind, err := r.aheadBehind(branch, upstream)
 	if err != nil {
 		return nil, err
 	}
 
 	bs.Upstream = upstream
-	bs.NeedsPush = needsPush
+	bs.Ahead = ahead
-	bs.NeedsPull = needsPull
+	bs.Behind = behind
 
 	return bs, nil
 }
@@ -215,61 +217,50 @@ func (r *Repo) upstream(branch string) (string, error) {
 	return remote + "/" + merge, nil
 }
 
-func (r *Repo) needsPullOrPush(localBranch string, upstreamBranch string) (needsPull bool, needsPush bool, err error) {
+func (r *Repo) aheadBehind(localBranch string, upstreamBranch string) (ahead int, behind int, err error) {
 	localHash, err := r.ResolveRevision(plumbing.Revision(localBranch))
 	if err != nil {
-		return false, false, errors.Wrapf(err, "Failed resolving revision %s", localBranch)
+		return 0, 0, errors.Wrapf(err, "Failed resolving revision %s", localBranch)
 	}
 
 	upstreamHash, err := r.ResolveRevision(plumbing.Revision(upstreamBranch))
 	if err != nil {
-		return false, false, errors.Wrapf(err, "Failed resolving revision %s", upstreamBranch)
+		return 0, 0, errors.Wrapf(err, "Failed resolving revision %s", upstreamBranch)
 	}
 
-	localCommit, err := r.CommitObject(*localHash)
+	behind, err = r.revlistCount(*localHash, *upstreamHash)
 	if err != nil {
-		return false, false, errors.Wrapf(err, "Failed finding a commit for hash %s", localHash.String())
+		return 0, 0, errors.Wrapf(err, "Failed counting commits behind %s", upstreamBranch)
 	}
 
-	upstreamCommit, err := r.CommitObject(*upstreamHash)
+	ahead, err = r.revlistCount(*upstreamHash, *localHash)
 	if err != nil {
-		return false, false, errors.Wrapf(err, "Failed finding a commit for hash %s", upstreamHash.String())
+		return 0, 0, errors.Wrapf(err, "Failed counting commits ahead of %s", upstreamBranch)
 	}
 
-	// If local branch hash is the same as upstream, it means there is no difference between local and upstream
+	return ahead, behind, nil
-	if *localHash == *upstreamHash {
+}
-		return false, false, nil
+
-	}
+// revlistCount counts the number of commits between two hashes.
-
+// https://github.com/src-d/go-git/issues/757#issuecomment-452697701
-	commons, err := localCommit.MergeBase(upstreamCommit)
+// TODO: See if this can be optimized. Running the loop twice feels wrong.
-	if err != nil {
+func (r *Repo) revlistCount(hash1, hash2 plumbing.Hash) (int, error) {
-		return false, false, errors.Wrapf(err, "Failed finding common ancestors for branches %s & %s", localBranch, upstreamBranch)
+	ref1hist, err := revlist.Objects(r.Storer, []plumbing.Hash{hash1}, nil)
-	}
+	if err != nil {
-
+		return 0, err
-	if len(commons) == 0 {
+	}
-		// TODO: No common ancestors. This should be an error
+
-		return false, false, nil
+	ref2hist, err := revlist.Objects(r.Storer, []plumbing.Hash{hash2}, ref1hist)
-	}
+	if err != nil {
-
+		return 0, err
-	if len(commons) > 1 {
+	}
-		// TODO: multiple best ancestors. How to handle this?
+
-		return false, false, nil
+	count := 0
-	}
+	for _, h := range ref2hist {
-
+		if _, err = r.CommitObject(h); err == nil {
-	mergeBase := commons[0]
+			count++
-
+		}
-	// If merge base is the same as upstream branch, local branch is ahead and push is needed
+	}
-	// If merge base is the same as local branch, local branch is behind and pull is needed
+
-	// If merge base is something else, branches have diverged and merge is needed (both pull and push)
+	return count, nil
-	// ref: https://stackoverflow.com/a/17723781/1085632
-
-	if mergeBase.Hash == *upstreamHash {
-		return false, true, nil
-	}
-
-	if mergeBase.Hash == *localHash {
-		return true, false, nil
-	}
-
-	return true, true, nil
 }

git/status_test.go

@@ -34,10 +34,10 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "master",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "",
+					Upstream: "",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				},
 			},
 		}},
@@ -47,10 +47,10 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "master",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "",
+					Upstream: "",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				},
 			},
 		}},
@@ -60,10 +60,10 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "master",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "",
+					Upstream: "",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				},
 			},
 		}},
@@ -73,15 +73,15 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "master",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "",
+					Upstream: "",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				}, {
-					Name:      "local",
+					Name:     "local",
-					Upstream:  "",
+					Upstream: "",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				},
 			},
 		}},
@@ -91,15 +91,15 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "local",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "origin/master",
+					Upstream: "origin/master",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				}, {
-					Name:      "local",
+					Name:     "local",
-					Upstream:  "",
+					Upstream: "",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				},
 			},
 		}},
@@ -109,10 +109,10 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         StatusDetached,
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "",
+					Upstream: "",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: false,
+					Ahead:    0,
 				},
 			},
 		}},
@@ -122,10 +122,10 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "master",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "origin/master",
+					Upstream: "origin/master",
-					NeedsPull: false,
+					Behind:   0,
-					NeedsPush: true,
+					Ahead:    1,
 				},
 			},
 		}},
@@ -135,10 +135,10 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "master",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "origin/master",
+					Upstream: "origin/master",
-					NeedsPull: true,
+					Behind:   1,
-					NeedsPush: false,
+					Ahead:    0,
 				},
 			},
 		}},
@@ -148,10 +148,10 @@ func TestStatus(t *testing.T) {
 			CurrentBranch:         "master",
 			Branches: []*BranchStatus{
 				{
-					Name:      "master",
+					Name:     "master",
-					Upstream:  "origin/master",
+					Upstream: "origin/master",
-					NeedsPull: true,
+					Behind:   3,
-					NeedsPush: true,
+					Ahead:    2,
 				},
 			},
 		}},

go.sum

@@ -31,6 +31,7 @@ github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
 github.com/bketelsen/crypt v0.0.3-0.20200106085610-5cbc8cc4026c/go.mod h1:MKsuJmJgSg28kpZDP6UIiPt0e0Oz0kqKNGyRaWEPv84=
+github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=
@@ -187,6 +188,7 @@ github.com/prometheus/common v0.0.0-20181113130724-41aa239b4cce/go.mod h1:daVV7q
 github.com/prometheus/common v0.4.0/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
 github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
 github.com/prometheus/procfs v0.0.0-20190507164030-5867b95ac084/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
+github.com/prometheus/tsdb v0.7.1 h1:YZcsG11NqnK4czYLrWd9mpEuAJIHVQLwdrleYfszMAA=
 github.com/prometheus/tsdb v0.7.1/go.mod h1:qhTCs0VvXwvX/y3TZrWD7rabWM+ijKTux40TwIPHuXU=
 github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=

print/print.go

@@ -67,14 +67,14 @@ func printBranchStatus(branch *git.BranchStatus) string {
 		status = append(status, fmt.Sprintf(ColorYellow, git.StatusNoUpstream))
 	}
 
-	if branch.NeedsPull {
+	if branch.Behind != 0 {
 		ok = false
-		status = append(status, fmt.Sprintf(ColorYellow, git.StatusBehind))
+		status = append(status, fmt.Sprintf(ColorYellow, fmt.Sprintf("%d %s", branch.Behind, git.StatusBehind)))
 	}
 
-	if branch.NeedsPush {
+	if branch.Ahead != 0 {
 		ok = false
-		status = append(status, fmt.Sprintf(ColorYellow, git.StatusAhead))
+		status = append(status, fmt.Sprintf(ColorYellow, fmt.Sprintf("%d %s", branch.Ahead, git.StatusAhead)))
 	}
 
 	if ok {