Drover's back, baby!

go/src/infra/libs/git/tree.go

Index: go/src/infra/libs/git/tree.go
diff --git a/go/src/infra/libs/git/tree.go b/go/src/infra/libs/git/tree.go
new file mode 100644
index 0000000000000000000000000000000000000000..9c9dc5b1b246dbf310e914c29d310a4f24aa85cb
--- /dev/null
+++ b/go/src/infra/libs/git/tree.go
@@ -0,0 +1,424 @@
+// 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 git
+
+import (
+	"bytes"
+	"fmt"
+	"sort"
+	"strconv"
+	"strings"
+
+	"infra/libs/infra_util"
+)
+
+// Contstants //////////////////////////////////////////////////////////////////
+
+const EmptyTreeHash = "4b825dc642cb6eb9a060e54bf8d69288fbee4904"
+
+// Types ///////////////////////////////////////////////////////////////////////
+
+// Tree is a true tree which represents a git 'tree' of objects, and can
+// represent those objects recursively (unlike a true git 'tree').
+//
+// Trees are also mutable, so you can manipulate them with impunity. Note that
+// ID() is not continuously updated, and you must call Rehash in order to
+// recalculate the correct ID() after manipulating the Tree.
+type Tree struct {
+	id              *ObjectID
+	cachedRawString *string
+
+	children map[string]*Child
+}
+
+// Constructors  ///////////////////////////////////////////////////////////////
+
+// NewEmptyTree creates a new Tree with the given |id|, but no children (e.g.
+// it is !Complete()). You only want to use this if you know that you're going
+// to be adding children to the Tree, otherwise use EmptyObject(TreeType, id).
+func NewEmptyTree(id Identifiable, childrenHint int) *Tree {
+	return &Tree{
+		id:       id.ID(),
+		children: make(map[string]*Child, childrenHint),
+	}
+}
+
+func BlankTreeChild(childrenHint int) *Child {
+	return &Child{NewEmptyTree(&NoID, childrenHint), 040000}
+}
+
+// TreeFromText returns a Complete() Tree for a `git ls-tree`-formatted
+// tree string.
+func NewTreeFromText(data []byte) (ret *Tree, err error) {
+	lines := strings.Split(strings.TrimSpace(string(data)), "\n")
+	retPtr := NewEmptyTree(&NoID, len(lines))
+	for _, line := range lines {
+		infoPath := strings.SplitN(line, "\t", 2)
+		if len(infoPath) != 2 {
+			return nil, fmt.Errorf("NewTreeFromText: unparsable line %#v", line)
+		}
+		var path string
+		if infoPath[1][0] == '"' {
+			path, err = strconv.Unquote(infoPath[1])
+			if err != nil {
+				return
+			}
+		} else {
+			path = infoPath[1]
+		}
+
+		modeTypeID := strings.Fields(infoPath[0])
+		if len(modeTypeID) != 3 {
+			return nil, fmt.Errorf("NewTreeFromText: unparsable line %#v", line)
+		}
+		var mode uint64
+		mode, err = strconv.ParseUint(modeTypeID[0], 8, 32)
+		if err != nil {
+			return
+		}
+
+		id, err := MakeObjectIDErr(modeTypeID[2])
+		if err != nil {
+			return nil, err
+		}
+		c, err := NewEmptyChild(Mode(mode), id)
+		if err != nil {
+			return nil, err
+		}
+		retPtr.SetChild(path, c)
+	}
+	retPtr.RawString()
+	ret = retPtr
+	return
+}
+
+func NewTreeFromRaw(data []byte) (ret *Tree, err error) {
+	return NewTreeFromRawWithID(MakeObjectIDForData(TreeType, data), data)
+}
+
+// TreeFromRawWithID trusts id, and will return a Complete() *Tree. Data is
+// formatted the same as RawString (i.e. in the raw git tree format).
+func NewTreeFromRawWithID(id Identifiable, data []byte) (ret *Tree, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = r.(error)
+		}
+	}()
+
+	retPtr := NewEmptyTree(id, 16)
+	// <mode> SP <name> NUL [20]byte
+	buf := bytes.NewBuffer(data)
+	nom := infra_util.Nom(buf)
+	yoink := infra_util.Yoink(buf)
+	for buf.Len() != 0 {
+		var mode uint64
+		mode, err = strconv.ParseUint(nom(' '), 8, 64)
+		if err != nil {
+			return
+		}
+		name := nom('\x00')
+
+		id := ObjectID{}
+		copy(id.raw[:], yoink(20))
+		c, err := NewEmptyChild(Mode(int(mode)), &id)
+		if err != nil {
+			return nil, err
+		}
+
+		retPtr.SetChild(name, c)
+	}
+	s := string(data)
+	retPtr.cachedRawString = &s
+	ret = retPtr
+	return
+}
+
+// Factories ///////////////////////////////////////////////////////////////////
+
+// LoadFullTree gets a recursively-enumerated Tree (or an EmptyObject)
+//
+// withBlobs:
+//   true  - Load blobs from Repo
+//   false - Blobs will be EmptyObject
+//
+// missingErr:
+//   true  - All entries in tree must load
+//   false - Missing entries will remain EmptyObject (or an !Complete() Tree)
+func LoadFullTree(s GitService, id Identifiable, withBlobs, missingErr bool) (Object, error) {
+	obj, err := s.GetObjectID(id)
+	if err != nil {
+		return nil, err
+	}
+	if obj == nil {
+		if !missingErr {
+			return NewEmptyObjectErr(TreeType, id)
+		}
+		return nil, fmt.Errorf("object %s missing", id)
+	}
+
+	switch x := obj.(type) {
+	case *Tree:
+	case *Commit:
+		info := s.ObjectInfo(x.ID().String() + ":")
+		if info == nil {
+			return nil, fmt.Errorf("object %s: missing", x.ID())
+		}
+		return LoadFullTree(s, info.ID, withBlobs, missingErr)
+	default:
+		return nil, fmt.Errorf("object %v is not a Tree!", x)
+	}
+	base := obj.(*Tree)
+
+	type change struct {
+		path  string
+		child *Child
+	}
+	changes := make(chan change, len(base.children))
+	allErrs := infra_util.FanOutIn(len(base.children), func(ch chan<- func() error) {
+		for p, c := range base.children {
+			p, c := p, c
+			ch <- func() (err error) {
+				var rslt Object
+
+				switch c.Mode.Type() {
+				case TreeType:
+					rslt, err = LoadFullTree(s, c.Object.ID(), withBlobs, missingErr)
+
+				case BlobType:
+					rslt = c.Object
+					if withBlobs {
+						rslt, err = s.GetObjectID(c.Object)
+					}
+				}
+
+				if err == nil {
+					if rslt != nil {
+						changes <- change{p, &Child{rslt, c.Mode}}
+					} else if missingErr {
+						err = fmt.Errorf("object %s missing", c.Object.ID())
+					}
+				}
+				return
+			}
+		}
+	})
+	close(changes)
+	for c := range changes {
+		base.children[c.path] = c.child
+	}
+
+	return base, allErrs
+}
+
+// Member functions ////////////////////////////////////////////////////////////
+
+func (t *Tree) String() string {
+	return fmt.Sprintf("Tree(%s, %v)", t.id, t.children)
+}
+
+// ID returns the last-calculated ObjectId of this tree. Note that this ID()
+// may be out-of-date if you have modified this Tree with {Set,Del}Child().
+func (t *Tree) ID() *ObjectID { return t.id }
+
+// Type returns TreeType to satisfy the Object interface
+func (t *Tree) Type() ObjectType { return TreeType }
+
+// Complete indicates that this tree's ID() matches its contents. It can
+// become false if you call DelChild or SetChild on this Tree.
+func (t *Tree) Complete() bool { return t.cachedRawString != nil }
+
+// NumChildren returns the number of immediate children this tree has.
+func (t *Tree) NumChildren() int { return len(t.children) }
+
+// Intern recursively interns this tree into the given repo.
+//
+// If the tree contains a missing object which is internable (i.g. not
+// EmptyObject), it will also be interned.
+//
+// If the tree contains an object with an ID of NoID, this method will panic.
+//
+// If the tree contains a missing object which is not internable, this method
+// will panic (see InternAllowMissing to negate this condition).
+//
+// This method has the side effect of causing all internable but !Complete()
+// objects to become Complete(). In particular, this means that the value of
+// t.ID() will be correct after calling this.
+func (t *Tree) Intern(s GitService, allowMissing bool) (*ObjectID, error) {
+	err := infra_util.FanOutIn(len(t.children), func(ch chan<- func() error) {
+		for _, c := range t.children {
+			c := c
+			ch <- func() (err error) {
+				switch x := c.Object.(type) {
+				case *Tree:
+					_, err = x.Intern(s, allowMissing)
+				case InternableObject:
+					s.Intern(x)
+				case Object:
+					if !s.HasObjectID(x) && !allowMissing {
+						return fmt.Errorf("missing non-internable object: %v", x)
+					}
+					if *x.ID() == NoID {
+						return fmt.Errorf("object has an ID of <NoID>")
+					}
+				}
+				return err
+			}
+		}
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return s.Intern(t)
+}
+
+// GetChild returns the Child at the slash-delemited |path| relative to this
+// Tree, or nil (if the child was not found).
+func (t *Tree) GetChild(path string) *Child {
+	pathParts := strings.SplitN(path, "/", 2)
+	if len(pathParts) == 1 {
+		return t.children[path]
+	}
+
+	dir, rest := pathParts[0], pathParts[1]
+	intermediate, exists := t.children[dir]
+	if !exists {
+		return nil
+	}
+
+	if childTree, ok := intermediate.Object.(*Tree); !ok {
+		panic(fmt.Errorf("entry path is not a tree: %s", dir))
+	} else {
+		return childTree.GetChild(rest)
+	}
+}
+
+// SetChild adds |child| to this Tree at the given slash-delimited |path|. It
+// will create intermediate empty Trees if necessary, but will panic() if you
+// attempt to set the child under something which is not a tree.
+func (t *Tree) SetChild(path string, child *Child) {
+	t.cachedRawString = nil
+	pathParts := strings.SplitN(path, "/", 2)
+	if len(pathParts) == 1 {
+		if len(path) == 0 {
+			panic(fmt.Errorf("attempting to SetChild with an empty path"))
+		}
+		t.children[path] = child
+	} else {
+		dir, rest := pathParts[0], pathParts[1]
+		intermediate, exists := t.children[dir]
+		if !exists {
+			intermediate = BlankTreeChild(1)
+			t.children[dir] = intermediate
+		}
+
+		switch x := intermediate.Object.(type) {
+		case *Tree:
+			x.SetChild(rest, child)
+			return
+		case *EmptyObject:
+			if x.Type() == TreeType {
+				if *x.ID() == NoID || x.ID().String() == EmptyTreeHash {
+					intermediate := BlankTreeChild(1)
+					t.children[dir] = intermediate
+					intermediate.Object.(*Tree).SetChild(rest, child)
+				} else {
+					panic(fmt.Errorf("entry path was a placeholder tree with an id: %v", x))
+				}
+				return
+			}
+		}
+		panic(fmt.Errorf("entry path is not a tree: %s", dir))
+	}
+}
+
+// DelChild removes the child at the slash-delimited |path|. Will panic if
+// |path| does not exist, or if one of the intermediate paths is not a Tree.
+func (t *Tree) DelChild(path string) bool {
+	t.cachedRawString = nil
+	pathParts := strings.SplitN(path, "/", 2)
+	if len(pathParts) == 1 {
+		var ret bool
+		_, ret = t.children[path]
+		delete(t.children, path)
+		return ret
+	} else {
+		dir, rest := pathParts[0], pathParts[1]
+		child, exists := t.children[dir]
+		if !exists {
+			return false
+		}
+
+		if childTree, ok := child.Object.(*Tree); !ok {
+			panic(fmt.Errorf("entry path is not a tree: %s", dir))
+		} else {
+			ret := childTree.DelChild(rest)
+			if childTree.NumChildren() == 0 {
+				delete(t.children, dir)
+			}
+			return ret
+		}
+	}
+}
+
+// RawString returns a git hash-object compatible representation of this Tree,
+// according to the currently-set children (e.g. it may return something which
+// does not match ID).
+func (t *Tree) RawString() string {
+	if t.cachedRawString == nil {
+		buf := &bytes.Buffer{}
+
+		entries := make(treeEntries, 0, len(t.children))
+		for path, child := range t.children {
+			switch x := child.Object.(type) {
+			case *Tree:
+				x.RawString()
+			}
+
+			entries = append(entries, treeEntry{
+				id:   child.Object.ID(),
+				mode: child.Mode,
+				name: path,
+			})
+		}
+
+		sort.Sort(entries)
+
+		for _, ent := range entries {
+			fmt.Fprintf(buf, "%o %s\x00%s", ent.mode, ent.name, ent.id.RawString())
+		}
+		s := buf.String()
+		t.id = MakeObjectIDForData(TreeType, buf.Bytes())
+		t.cachedRawString = &s
+	}
+	return *t.cachedRawString
+}
+
+/// Private ///////////////////////////////////////////////////////////////////
+
+type treeEntry struct {
+	id   *ObjectID
+	mode Mode
+	name string
+}
+type treeEntries []treeEntry
+
+func (s treeEntries) Len() int { return len(s) }
+
+// Extra trailing slash is how git sorts the tree objects too.
+func (s treeEntries) Less(i, j int) bool {
+	a, b := s[i].name, s[j].name
+
+	if s[i].mode.Type() == TreeType {
+		a += "/"
+	}
+	if s[j].mode.Type() == TreeType {
+		b += "/"
+	}
+
+	return a < b
+}
+func (s treeEntries) Swap(i, j int) { s[i], s[j] = s[j], s[i] }