Drover's back, baby!

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

+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
+}