1 files changed, 65 insertions, 49 deletions

diff --git a/tree.go b/tree.go
index 4e67f0e..9a750e0 100644
--- a/tree.go
+++ b/tree.go
@@ -5,18 +5,28 @@ import (
 	"errors"
 	"io"
 	"os"
-	"path"
 	"strconv"
 	"strings"
 
 	"gopkg.in/src-d/go-git.v3/core"
 )
 
+const (
+	maxTreeDepth = 1024
+)
+
+// New errors defined by this package.
+var (
+	ErrMaxTreeDepth = errors.New("maximum tree depth exceeded")
+	ErrFileNotFound = errors.New("file not found")
+)
+
 // Tree is basically like a directory - it references a bunch of other trees
 // and/or blobs (i.e. files and sub-directories)
 type Tree struct {
-	Entries map[string]TreeEntry
-	Hash    core.Hash
+	Entries      map[string]TreeEntry
+	OrderedNames []string
+	Hash         core.Hash
 
 	r *Repository
 }
@@ -28,9 +38,6 @@ type TreeEntry struct {
 	Hash core.Hash
 }
 
-// New errors defined by this package.
-var ErrFileNotFound = errors.New("file not found")
-
 // File returns the hash of the file identified by the `path` argument.
 // The path is interpreted as relative to the tree receiver.
 func (t *Tree) File(path string) (*File, error) {
@@ -113,40 +120,8 @@ func (t *Tree) entry(baseName string) (*TreeEntry, error) {
 	return &entry, nil
 }
 
-func (t *Tree) Files() chan *File {
-	ch := make(chan *File, 1)
-
-	go func() {
-		defer func() { close(ch) }()
-		t.walkEntries("", ch)
-	}()
-
-	return ch
-}
-
-func (t *Tree) walkEntries(base string, ch chan *File) {
-	for _, entry := range t.Entries {
-		obj, err := t.r.Storage.Get(entry.Hash)
-		if err != nil {
-			if err == core.ObjectNotFoundErr {
-				continue // ignore entries without hash (= submodule dirs)
-			}
-			//FIXME: Refactor this function to return an error. Ideally this would be
-			//       moved into a FileIter type.
-		}
-
-		if obj.Type() == core.TreeObject {
-			tree := &Tree{r: t.r}
-			tree.Decode(obj)
-			tree.walkEntries(path.Join(base, entry.Name), ch)
-			continue
-		}
-
-		blob := &Blob{}
-		blob.Decode(obj)
-
-		ch <- &File{Name: path.Join(base, entry.Name), Reader: blob.Reader(), Hash: entry.Hash}
-	}
+func (t *Tree) Files() *FileIter {
+	return NewFileIter(t.r, t)
 }
 
 // Decode transform an core.Object into a Tree struct
@@ -195,26 +170,67 @@ func (t *Tree) Decode(o core.Object) error {
 			Mode: os.FileMode(fm),
 			Name: baseName,
 		}
+		t.OrderedNames = append(t.OrderedNames, baseName)
 	}
 
 	return nil
 }
 
+// TreeEntryIter facilitates iterating through the TreeEntry objects in a Tree.
+type TreeEntryIter struct {
+	t   *Tree
+	pos int
+}
+
+func NewTreeEntryIter(t *Tree) *TreeEntryIter {
+	return &TreeEntryIter{t, 0}
+}
+
+func (iter *TreeEntryIter) Next() (TreeEntry, error) {
+	if iter.pos >= len(iter.t.OrderedNames) {
+		return TreeEntry{}, io.EOF
+	}
+
+	entry, ok := iter.t.Entries[iter.t.OrderedNames[iter.pos]]
+	if !ok {
+		// Probable race condition or internal bug
+		// FIXME: Report more severe error or panic
+		return TreeEntry{}, io.EOF
+	}
+
+	iter.pos++
+
+	return entry, nil
+}
+
+// TreeEntryIter facilitates iterating through the descendent subtrees of a
+// Tree.
 type TreeIter struct {
-	core.ObjectIter
-	r *Repository
+	w TreeWalker
 }
 
-func NewTreeIter(r *Repository, iter core.ObjectIter) *TreeIter {
-	return &TreeIter{iter, r}
+func NewTreeIter(r *Repository, t *Tree) *TreeIter {
+	return &TreeIter{
+		w: *NewTreeWalker(r, t),
+	}
 }
 
 func (iter *TreeIter) Next() (*Tree, error) {
-	obj, err := iter.ObjectIter.Next()
-	if err != nil {
-		return nil, err
+	for {
+		_, _, obj, err := iter.w.Next()
+		if err != nil {
+			return nil, err
+		}
+
+		if obj.Type() != core.TreeObject {
+			// Skip non-tree objects
+			continue
+		}
+
+		return iter.w.Tree(), nil
 	}
+}
 
-	tree := &Tree{r: iter.r}
-	return tree, tree.Decode(obj)
+func (iter *TreeIter) Close() {
+	iter.w.Close()
 }