Implement streaming Tree reader

tree.go

@@ -2,7 +2,6 @@ package git
 
 import (
 	"errors"
-	"io/ioutil"
 	"path"
 	"strings"
 )
@@ -99,26 +98,22 @@ func (t *Tree) ListEntries() Entries {
 
 	t.entriesParsed = true
 
-	_, _, dataRc, err := t.repo.getRawObject(t.Id, false)
+	var entries Entries
+
+	scanner, err := t.Scanner()
 	if err != nil {
 		return nil
 	}
 
-	defer func() {
-		dataRc.Close()
-	}()
-
-	// TODO reader
-	data, err := ioutil.ReadAll(dataRc)
-	if err != nil {
-		return nil
+	for scanner.Scan() {
+		entries = append(entries, scanner.TreeEntry())
 	}
 
-	t.entries, err = parseTreeData(t, data)
-	if err != nil {
+	if err := scanner.Err(); err != nil {
 		return nil
 	}
 
+	t.entries = entries
 	return t.entries
 }
 
@@ -128,3 +123,11 @@ func NewTree(repo *Repository, id sha1) *Tree {
 	tree.repo = repo
 	return tree
 }
+
+func (t *Tree) Scanner() (*TreeScanner, error) {
+	_, _, r, err := t.repo.getRawObject(t.Id, false)
+	if err != nil {
+		return nil, err
+	}
+	return NewTreeScanner(t.Id, r), nil
+}

tree_utils.go

@@ -1,50 +1,135 @@
 package git
 
 import (
+	"bufio"
 	"bytes"
-	"errors"
+	"fmt"
+	"io"
+	"regexp"
 )
 
-// Parse tree information from the (uncompressed) raw
-// data from the tree object.
-func parseTreeData(tree *Tree, data []byte) ([]*TreeEntry, error) {
-	entries := make([]*TreeEntry, 0, 10)
-	l := len(data)
-	pos := 0
-	for pos < l {
-		entry := new(TreeEntry)
-		entry.ptree = tree
-		spacepos := bytes.IndexByte(data[pos:], ' ')
-		switch string(data[pos : pos+spacepos]) {
-		case "100644":
-			entry.mode = ModeBlob
-			entry.Type = ObjectBlob
-		case "100755":
-			entry.mode = ModeExec
-			entry.Type = ObjectBlob
-		case "120000":
-			entry.mode = ModeSymlink
-			entry.Type = ObjectBlob
-		case "160000":
-			entry.mode = ModeCommit
-			entry.Type = ObjectCommit
-		case "40000":
-			entry.mode = ModeTree
-			entry.Type = ObjectTree
-		default:
-			return nil, errors.New("unknown type: " + string(data[pos:pos+spacepos]))
-		}
-		pos += spacepos + 1
-		zero := bytes.IndexByte(data[pos:], 0)
-		entry.name = string(data[pos : pos+zero])
-		pos += zero + 1
-		id, err := NewId(data[pos : pos+20])
-		if err != nil {
-			return nil, err
+type TreeScanner struct {
+	parent sha1
+
+	*bufio.Scanner
+	closer io.Closer
+
+	treeEntry *TreeEntry
+	err       error
+}
+
+func NewTreeScanner(parent sha1, rc io.ReadCloser) *TreeScanner {
+	ts := &TreeScanner{
+		parent:  parent,
+		Scanner: bufio.NewScanner(rc),
+		closer:  rc,
+	}
+	ts.Split(ScanTreeEntry)
+	return ts
+}
+
+var TreeEntryRe = regexp.MustCompile("^([0-9]+) ([^\x00]+)\x00")
+
+func (t *TreeScanner) parse() error {
+	t.treeEntry = nil
+
+	data := t.Bytes()
+
+	match := TreeEntryRe.FindSubmatch(data)
+	if match == nil {
+		return fmt.Errorf("failed to parse tree entry: %q", data)
+	}
+
+	modeString, name := string(match[1]), string(match[2])
+	id, err := NewId(data[len(match[0]):])
+	if err != nil {
+		return err
+	}
+
+	entryMode, objectType, err := ParseModeType(modeString)
+	if err != nil {
+		return err
+	}
+
+	t.treeEntry = &TreeEntry{
+		name: name,
+		mode: entryMode,
+		Id:   id,
+		Type: objectType,
+	}
+
+	return nil
+}
+
+func (t *TreeScanner) Scan() bool {
+	if !t.Scanner.Scan() {
+		if t.closer != nil {
+			// Upon hitting any error, close the input.
+			t.closer.Close()
+			t.closer = nil
 		}
-		entry.Id = id
-		pos = pos + 20
-		entries = append(entries, entry)
+		return false
+	}
+
+	t.err = t.parse()
+	if t.err != nil {
+		return false
+	}
+
+	return true
+}
+
+func (t *TreeScanner) Err() error {
+	// Underlying IO errs take priority
+	if err := t.Scanner.Err(); err != nil {
+		return err
+	}
+	return t.err
+}
+
+func ScanTreeEntry(
+	data []byte,
+	atEOF bool,
+) (
+	advance int, token []byte, err error,
+) {
+	if atEOF && len(data) == 0 {
+		return 0, nil, nil
+	}
+	const shaLen = 20
+
+	nullIndex := bytes.IndexByte(data, '\x00')
+	recordLength := nullIndex + 1 + shaLen
+	if recordLength <= len(data) {
+		// We found 20 bytes after a null, we're done.
+		return recordLength, data[:recordLength], nil
+	}
+
+	if atEOF {
+		// atEOF but don't have a complete record
+		return 0, nil, fmt.Errorf("malformed record %q", data)
+	}
+
+	return 0, nil, nil // Request more data.
+}
+
+func (t *TreeScanner) TreeEntry() *TreeEntry {
+	return t.treeEntry
+}
+
+func ParseModeType(modeString string) (EntryMode, ObjectType, error) {
+	switch modeString {
+	case "100644":
+		return ModeBlob, ObjectBlob, nil
+	case "100755":
+		return ModeExec, ObjectBlob, nil
+	case "120000":
+		return ModeSymlink, ObjectBlob, nil
+	case "160000":
+		return ModeCommit, ObjectCommit, nil
+	case "40000":
+		return ModeTree, ObjectTree, nil
+	default:
 	}
-	return entries, nil
+	return 0, 0, fmt.Errorf("unknown type: %q", modeString)
 }