Drover's back, baby!

go/src/infra/tools/drover/merge.go

Index: go/src/infra/tools/drover/merge.go
diff --git a/go/src/infra/tools/drover/merge.go b/go/src/infra/tools/drover/merge.go
new file mode 100644
index 0000000000000000000000000000000000000000..271e1a6c9abf7381fa16eb2f08266f0d60f2533d
--- /dev/null
+++ b/go/src/infra/tools/drover/merge.go
@@ -0,0 +1,365 @@
+package main
+
+import "bytes"
+import "fmt"
+import "strings"
+import "sync"
+
+import "github.com/cheggaaa/pb"
+import "github.com/daviddengcn/go-colortext"
+
+import "infra/libs/git"
+import "infra/libs/gitiles"
+
+type stringSet map[string]struct{}
+
+func (s stringSet) Has(k string) bool {
+	_, ok := s[k]
+	return ok
+}
+
+func (s stringSet) Add(k string) bool {
+	if s.Has(k) {
+		return false
+	}
+	s[k] = struct{}{}
+	return true
+}
+
+func (s stringSet) Del(k string) {
+	delete(s, k)
+}
+
+func (s stringSet) String() string {
+	buf := &bytes.Buffer{}
+	fmt.Fprint(buf, "{")
+	first := true
+	for k := range s {
+		if !first {
+			fmt.Fprint(buf, " ")
+		}
+		first = false
+		fmt.Fprint(buf, k)
+	}
+	fmt.Fprint(buf, "}")
+	return buf.String()
+}
+
+type internRequest struct {
+	commit     git.ObjectID
+	pathPieces []string
+}
+
+func (i internRequest) repoName() (ret string) {
+	ret = i.commit.String()
+	if len(i.pathPieces) > 0 {
+		ret += ":" // root tree
+		ret += strings.Join(i.pathPieces, "/")
+	}
+	return
+}
+
+func commitRequest(commit git.ObjectID) internRequest {
+	return internRequest{commit: commit}
+}
+
+func rootTreeRequest(commit git.ObjectID) internRequest {
+	return internRequest{commit: commit, pathPieces: []string{""}}
+}
+
+func objRequest(commit git.ObjectID, pathPieces []string) internRequest {
+	for _, p := range pathPieces {
+		if p == "" {
+			panic(fmt.Errorf(
+				"Cannot have an empty pathPieces entry in an objRequest! %#v", pathPieces))
+		}
+	}
+	return internRequest{commit: commit, pathPieces: pathPieces}
+}
+
+type progMessage bool
+
+const (
+	progDone progMessage = iota == 0
+	progAdd
+)
+
+type lazyProgBar chan<- progMessage
+
+func newLazyProgBar(estSize int) lazyProgBar {
+	ch := make(chan progMessage, estSize)
+	var bar *pb.ProgressBar
+	go func() {
+		for amt := range ch {
+			if bar == nil {
+				if amt == progDone {
+					panic("cannot start progress with a progDone")
+				}
+				bar = pb.StartNew(1)
+			} else {
+				switch amt {
+				case progAdd:
+					bar.Total++
+				case progDone:
+					bar.Increment()
+				}
+			}
+		}
+		if bar != nil {
+			bar.Finish()
+		}
+	}()
+	return ch
+}
+
+func (l lazyProgBar) Add()    { l <- progAdd }
+func (l lazyProgBar) Done()   { l <- progDone }
+func (l lazyProgBar) Finish() { close(l) }
+
+func internService(r *git.Repo, g *gitiles.Gitiles, reqs []internRequest) <-chan error {
+	bar := newLazyProgBar(len(reqs))
+	grp := sync.WaitGroup{}
+	grp.Add(len(reqs))
+	allErrCh := make(chan error, len(reqs))
+
+	go func() {
+		grp.Wait()
+		bar.Finish()
+		close(allErrCh)
+	}()
+
+	for _, req := range reqs {
+		req := req
+		go func() {
+			defer grp.Done()
+			if r.HasObject(req.repoName()) {
+				return
+			}
+			bar.Add()
+			defer bar.Done()
+			rslt := <-g.GetObjectFromPath(req.commit.String(), req.pathPieces...)
+			err := rslt.Err
+			if err == nil {
+				irslt := r.Intern(rslt.Object)
+				err = irslt.Err
+			}
+			allErrCh <- err
+		}()
+	}
+
+	return allErrCh
+}
+
+func acquireObjects(r *git.Repo, g *gitiles.Gitiles, commit, landCommit *git.Commit, treeDiff git.TreeDiff) {
+	if len(commit.Parents()) != 1 {
+		panic(fmt.Errorf("Got wrong number of parents for commit %s: %s",
+			commit.ID(), commit.Parents()))
+	}
+
+	fmt.Println("Acquiring objects")
+	reqs := make([]internRequest, 0, 5*3*len(treeDiff))
+
+	dedup := stringSet{}
+
+	add := func(req internRequest) {
+		if !dedup.Add(req.repoName()) {
+			return
+		}
+		reqs = append(reqs, req)
+	}
+
+	addBlobTreesFor := func(commit git.ObjectID, path string) {
+		if path == "/dev/null" {
+			return
+		}
+		pathBits := strings.Split(path, "/")
+		add(objRequest(commit, pathBits))
+		add(rootTreeRequest(commit))
+		for i := 1; i < len(pathBits); i++ {
+			add(objRequest(commit, pathBits[:i]))
+		}
+	}
+
+	add(commitRequest(commit.ID()))
+	add(commitRequest(commit.Parents()[0]))
+	add(commitRequest(landCommit.ID()))
+	for _, wholeEntry := range treeDiff {
+		addBlobTreesFor(commit.ID(), wholeEntry.New.Name)
+		addBlobTreesFor(commit.Parents()[0], wholeEntry.Old.Name)
+		addBlobTreesFor(landCommit.ID(), wholeEntry.Old.Name)
+	}
+
+	for e := range internService(r, g, reqs) {
+		failIf(e)
+	}
+	fmt.Println("Done")
+}
+
+func hollow(r *git.Repo, commit git.ObjectID, parents []git.ObjectID, t *git.Tree) git.ObjectID {
+	t.Intern(r)
+	cur := r.GetObjectID(commit).(*git.Commit)
+	if cur == nil {
+		panic("Could not find commit: " + commit.String())
+	}
+	h := cur.SetTree(t.ID()).SetParents(parents)
+	rslt := r.Intern(&h)
+	if rslt.Err != nil {
+		panic(rslt.Err)
+	}
+	return rslt.ID
+}
+
+func resolveConflict(r *git.Repo) []string {
+	conflicts := []string{}
+	for _, line := range SplitLines(r.RunOutput("status", "--porcelain")) {
+		if strings.Contains(line[:2], "U") {
+			conflicts = append(conflicts, line[3:])
+		}
+	}
+
+	ct.ChangeColor(ct.Red, false, ct.None, false)
+	fmt.Println()
+	fmt.Println("There was a conflict during the cherry-pick operation.")
+	fmt.Println("Please resolve it. If you exit the shell without resolving")
+	fmt.Println("it, drover will abort.")
+	fmt.Println()
+	ct.ResetColor()
+	fmt.Println("Recap:")
+	ct.ChangeColor(ct.Green, false, ct.None, false)
+	fmt.Println("  until no conflicted files:")
+	fmt.Println("    git status                # See conflicted files")
+	fmt.Println("    $EDITOR <file>            # Edit conflicted files")
+	fmt.Println("    git add <file>            # Let cherry pick know you're done with <file>")
+	ct.ChangeColor(ct.Cyan, false, ct.None, false)
+	fmt.Println("  git cherry-pick --continue  # Message doesn't matter")
+	fmt.Println("  exit                        # Resume drover")
+	fmt.Println()
+	ct.ResetColor()
+	fmt.Println("Conflicts:")
+	ct.ChangeColor(ct.Red, false, ct.None, false)
+	for _, c := range conflicts {
+		fmt.Println("  " + c)
+	}
+	ct.ResetColor()
+	fmt.Println()
+
+	Shell(r.Path)
+
+	if r.RunOk("rev-parse", "CHERRY_PICK_HEAD") {
+		r.Run("cherry-pick", "--abort")
+		failIf(fmt.Errorf("Aborting due to unresolved conflict."))
+	}
+
+	return conflicts
+}
+
+func getNewMessage(commit *git.Commit, mode string, conflicts []string) string {
+	var msgLines []string
+	add := func(lines ...string) { msgLines = append(msgLines, lines...) }
+	addConflicts := func() {
+		if len(conflicts) != 0 {
+			add("", "Conflicts:")
+			for _, c := range conflicts {
+				add("\t" + c)
+			}
+		}
+	}
+	if mode == "revert" {
+		rawLines := commit.MessageRawLines()
+		msgLines = make([]string, 0, len(rawLines)+2)
+		add("Revert \"" + rawLines[0] + "\"")
+		add("")
+		add("This reverts commit " + commit.ID().String() + ".")
+		add("")
+		add("Original commit message:")
+		for _, l := range rawLines {
+			add("> " + l)
+		}
+		addConflicts()
+	} else {
+		msgLines = commit.MessageLines()
+		addConflicts()
+		add("")
+		for _, f := range commit.FooterPairs() {
+			add(fmt.Sprintf("%s: %s", f.Key, f.Value))
+		}
+		add(fmt.Sprintf("(cherry picked from commit %s)\n", commit.ID()))
+	}
+	return strings.Join(msgLines, "\n")
+}
+
+func createCommit(mode string, r *git.Repo, parent, commit, land git.ObjectID) *git.Commit {
+	addChild := func(t *git.Tree, e git.TreeDiffEntryHalf) {
+		switch e.Mode.Type() {
+		case "tree":
+			return
+		case "blob":
+			c := e.Child
+			t.SetChild(e.Name, &c)
+		default:
+			panic(fmt.Sprintf("Cannot process mode: %d", e.Mode))
+		}
+	}
+
+	diffEnts := r.DiffTree(parent.String(), commit.String())
+
+	parTree := git.NewEmptyTree(git.NoID, len(diffEnts))
+	landTree := git.NewEmptyTree(git.NoID, len(diffEnts))
+	cmtTree := git.NewEmptyTree(git.NoID, len(diffEnts))
+
+	for _, e := range diffEnts {
+		addChild(parTree, e.Old)
+		addChild(cmtTree, e.New)
+
+		landID, ok := r.Run("rev-parse", land.String()+":"+e.Old.Name)
+		if !ok {
+			panic(fmt.Errorf("Could not rev-parse landing blob: %s:%s", land, e.Old.Name))
+		}
+		addChild(landTree, git.TreeDiffEntryHalf{
+			Name: e.Old.Name,
+			Child: *git.NewEmptyChild(
+				e.Old.Mode, git.MakeObjectID(strings.TrimSpace(landID))),
+		})
+	}
+
+	parHollow := hollow(r, parent, []git.ObjectID{}, parTree)
+	landHollow := hollow(r, land, []git.ObjectID{parHollow}, landTree)
+	cmtHollow := hollow(r, commit, []git.ObjectID{parHollow}, cmtTree)
+
+	conflicts := []string{}
+
+	r.Run("reset", "--hard", landHollow.String())
+	r.Run("checkout", "-f", landHollow.String())
+	if !r.RunOk("cherry-pick", "-Xpatience", cmtHollow.String()) {
+		conflicts = resolveConflict(r)
+	}
+
+	// TODO(riannucci): upload for review if len(conflicts) != 0
+
+	fullTree := r.GetFullTree(land.String()+":", git.NoBlobs, git.MissingOK)
+	for _, e := range r.DiffTree(landHollow.String(), "HEAD") {
+		if e.Old.Mode.Type() == "blob" {
+			fullTree.DelChild(e.Old.Name)
+		}
+		if e.New.Mode.Type() == "blob" {
+			c := e.New.Child
+			fullTree.SetChild(e.New.Name, &c)
+		}
+	}
+	fullTree.InternAllowMissing(r)
+
+	cmtMsgSrc := commit
+	if mode == "revert" {
+		cmtMsgSrc = parent
+	}
+	msg := getNewMessage(r.GetObjectID(cmtMsgSrc).(*git.Commit), mode, conflicts)
+
+	cpick := (r.GetObject("HEAD").(*git.Commit).
+		SetTree(fullTree.ID()).
+		SetParents([]git.ObjectID{land}).
+		SetRawMessage(msg))
+	rslt := r.Intern(&cpick)
+	failIf(rslt.Err)
+
+	return &cpick
+}