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..94d6fbf74d79a25361cefb0e874089b2c331c150
--- /dev/null
+++ b/go/src/infra/tools/drover/merge.go
@@ -0,0 +1,366 @@
+// Copyright 2014 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+
+	"github.com/cheggaaa/pb"
+	"github.com/daviddengcn/go-colortext"
+
+	"infra/libs/git"
+	"infra/libs/git/repo"
+	"infra/libs/gitiles"
+	"infra/libs/infra_util"
+)
+
+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 *repo.Repo, g *gitiles.Gitiles, reqs []internRequest) error {
+	bar := newLazyProgBar(len(reqs))
+	defer bar.Finish()
+
+	return infra_util.FanOutIn(len(reqs), func(ch chan<- func() error) {
+		for _, req := range reqs {
+			req := req
+			ch <- func() error {
+				if r.HasObject(req.repoName()) {
+					return nil
+				}
+				bar.Add()
+				defer bar.Done()
+				rslt, err := g.GetObjectFromPath(req.commit.String(), req.pathPieces...)
+				if err == nil {
+					_, err = r.Intern(rslt)
+				}
+				return err
+			}
+		}
+	})
+}
+
+func acquireObjects(r *repo.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)
+	}
+
+	err := internService(r, g, reqs)
+	failIf(err)
+	fmt.Println("Done")
+}
+
+func hollow(r *repo.Repo, commit *git.ObjectID, parents []*git.ObjectID, t *git.Tree) *git.ObjectID {
+	t.Intern(r, false)
+	cur_obj, err := r.GetObjectID(commit)
+	failIf(err)
+
+	rslt, err := r.Intern(cur_obj.(*git.Commit).SetTree(t.ID()).SetParents(parents))
+	if err != nil {
+		panic(err)
+	}
+	return rslt
+}
+
+func resolveConflict(r *repo.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.MustRunOk("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 *repo.Repo, parent, commit, land *git.ObjectID) *git.Commit {
+	addChild := func(t *git.Tree, e git.TreeDiffEntryHalf) {
+		switch e.Mode.Type() {
+		case git.TreeType:
+			return
+		case git.BlobType:
+			c := e.Child
+			t.SetChild(e.Name, &c)
+		default:
+			panic(fmt.Sprintf("Cannot process mode: %d", e.Mode))
+		}
+	}
+
+	diffEnts, err := r.DiffTree(parent.String(), commit.String())
+	failIf(err)
+
+	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 := r.MustRunOutput("rev-parse", land.String()+":"+e.Old.Name)
+		id := git.MakeObjectID(strings.TrimSpace(landID))
+		ch, err := git.NewEmptyChild(e.Old.Mode, id)
+		failIf(err)
+
+		addChild(landTree, git.TreeDiffEntryHalf{
+			Name:  e.Old.Name,
+			Child: *ch,
+		})
+	}
+
+	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.MustRunOk("reset", "--hard", landHollow.String())
+	r.MustRunOk("checkout", "-f", landHollow.String())
+	if !r.RunOk("cherry-pick", "-Xpatience", cmtHollow.String()) {
+		conflicts = resolveConflict(r)
+	}
+
+	// TODO(iannucci): upload for review if len(conflicts) != 0
+
+	fullTreeObj, err := git.LoadFullTree(r, land, false, false)
+	failIf(err)
+	fullTree := fullTreeObj.(*git.Tree)
+	diff, err := r.DiffTree(landHollow.String(), "HEAD")
+	failIf(err)
+	for _, e := range diff {
+		if e.Old.Mode.Type() == git.BlobType {
+			fullTree.DelChild(e.Old.Name)
+		}
+		if e.New.Mode.Type() == git.BlobType {
+			c := e.New.Child
+			fullTree.SetChild(e.New.Name, &c)
+		}
+	}
+	fullTree.Intern(r, true)
+
+	cmtMsgSrc := commit
+	if mode == "revert" {
+		cmtMsgSrc = parent
+	}
+	cmtObj, err := r.GetObjectID(cmtMsgSrc)
+	failIf(err)
+	msg := getNewMessage(cmtObj.(*git.Commit), mode, conflicts)
+
+	cmtObj, err = r.GetObject("HEAD")
+	failIf(err)
+	cpick := (cmtObj.(*git.Commit).
+		SetTree(fullTree.ID()).
+		SetParents([]*git.ObjectID{land}).
+		SetRawMessage(msg))
+	_, err = r.Intern(cpick)
+	failIf(err)
+
+	return cpick
+}