TreeWalker optimization

4 files changed, 35 insertions, 16 deletions

diff --git a/file.go b/file.go
index ef17c6e..e18caaf 100644
--- a/file.go
+++ b/file.go
@@ -53,23 +53,31 @@ func (f *File) Lines() ([]string, error) {
 }
 
 type FileIter struct {
+	r *Repository
 	w TreeWalker
 }
 
 func NewFileIter(r *Repository, t *Tree) *FileIter {
-	return &FileIter{w: *NewTreeWalker(r, t)}
+	return &FileIter{r: r, w: *NewTreeWalker(r, t)}
 }
 
 func (iter *FileIter) Next() (*File, error) {
 	for {
-		name, entry, obj, err := iter.w.Next()
+		name, entry, err := iter.w.Next()
 		if err != nil {
 			return nil, err
 		}
 
-		if blob, ok := obj.(*Blob); ok {
-			return newFile(name, entry.Mode, blob), nil
+		if entry.Mode.IsDir() {
+			continue
 		}
+
+		blob, err := iter.r.Blob(entry.Hash)
+		if err != nil {
+			return nil, err
+		}
+
+		return newFile(name, entry.Mode, blob), nil
 	}
 }
 
diff --git a/tree.go b/tree.go
index 95fbcb5..eaa0c15 100644
--- a/tree.go
+++ b/tree.go
@@ -269,12 +269,14 @@ func (iter *treeEntryIter) Next() (TreeEntry, error) {
 // TreeEntryIter facilitates iterating through the descendent subtrees of a
 // Tree.
 type TreeIter struct {
+	r *Repository
 	w TreeWalker
 }
 
 // NewTreeIter returns a new TreeIter instance
 func NewTreeIter(r *Repository, t *Tree) *TreeIter {
 	return &TreeIter{
+		r: r,
 		w: *NewTreeWalker(r, t),
 	}
 }
@@ -282,14 +284,16 @@ func NewTreeIter(r *Repository, t *Tree) *TreeIter {
 // Next returns the next Tree from the tree.
 func (iter *TreeIter) Next() (*Tree, error) {
 	for {
-		_, _, obj, err := iter.w.Next()
+		_, entry, err := iter.w.Next()
 		if err != nil {
 			return nil, err
 		}
 
-		if tree, ok := obj.(*Tree); ok {
-			return tree, nil
+		if !entry.Mode.IsDir() {
+			continue
 		}
+
+		return iter.r.Tree(entry.Hash)
 	}
 }
 
diff --git a/tree_walker.go b/tree_walker.go
index d4aa01a..cbd81c4 100644
--- a/tree_walker.go
+++ b/tree_walker.go
@@ -43,7 +43,8 @@ func NewTreeWalker(r *Repository, t *Tree) *TreeWalker {
 // In the current implementation any objects which cannot be found in the
 // underlying repository will be skipped automatically. It is possible that this
 // may change in future versions.
-func (w *TreeWalker) Next() (name string, entry TreeEntry, obj Object, err error) {
+func (w *TreeWalker) Next() (name string, entry TreeEntry, err error) {
+	var obj Object
 	for {
 		current := len(w.stack) - 1
 		if current < 0 {
@@ -51,6 +52,7 @@ func (w *TreeWalker) Next() (name string, entry TreeEntry, obj Object, err error
 			err = io.EOF
 			return
 		}
+
 		if current > maxTreeDepth {
 			// We're probably following bad data or some self-referencing tree
 			err = ErrMaxTreeDepth
@@ -65,6 +67,7 @@ func (w *TreeWalker) Next() (name string, entry TreeEntry, obj Object, err error
 			w.base = path.Clean(w.base) // Remove trailing slash
 			continue
 		}
+
 		if err != nil {
 			return
 		}
@@ -73,10 +76,9 @@ func (w *TreeWalker) Next() (name string, entry TreeEntry, obj Object, err error
 			err = nil
 			continue
 		}
+
 		if entry.Mode.IsDir() {
 			obj, err = w.r.Tree(entry.Hash)
-		} else {
-			obj, err = w.r.Blob(entry.Hash)
 		}
 
 		name = path.Join(w.base, entry.Name)
diff --git a/tree_walker_test.go b/tree_walker_test.go
index c03ddb2..fe6bcae 100644
--- a/tree_walker_test.go
+++ b/tree_walker_test.go
@@ -104,15 +104,20 @@ func (s *SuiteTreeWalker) TestNext(c *C) {
 		for k := 0; k < len(t.objs); k++ {
 			info := t.objs[k]
 			c.Assert(err, IsNil)
-			name, entry, obj, err := walker.Next()
+			name, entry, err := walker.Next()
 
 			c.Assert(err, IsNil, Commentf("subtest %d, iter %d, err=%v", i, k, err))
-			c.Assert(name, Equals, info.Name, Commentf("subtest %d, iter %d, name=%v, expected=%s, stack=%v, base=%v", i, k, name, info.Name, walker.stack, walker.base))
-			c.Assert(obj.Type(), Equals, info.Kind, Commentf("subtest %d, iter %d, obj.Type()=%v expected=%v", i, k, obj.Type(), info.Kind))
-			c.Assert(entry.Hash.String(), Equals, info.Hash, Commentf("subtest %d, iter %d, entry.Hash=%v, expected=%s", i, k, entry.Hash, info.Hash))
-			c.Assert(obj.ID().String(), Equals, info.Hash, Commentf("subtest %d, iter %d, obj.ID()=%v, expected=%s", i, k, obj.ID(), info.Hash))
+			c.Assert(name, Equals, info.Name,
+				Commentf("subtest %d, iter %d, name=%v, expected=%s, stack=%v, base=%v", i, k, name, info.Name, walker.stack, walker.base))
+
+			c.Assert(entry.Hash.String(), Equals, info.Hash,
+				Commentf("subtest %d, iter %d, entry.Hash=%v, expected=%s", i, k, entry.Hash, info.Hash))
+
+			c.Assert(entry.Hash.String(), Equals, info.Hash,
+				Commentf("subtest %d, iter %d, obj.ID()=%v, expected=%s", i, k, entry.Hash.String(), info.Hash))
 		}
-		_, _, _, err = walker.Next()
+
+		_, _, err = walker.Next()
 		c.Assert(err, Equals, io.EOF)
 	}
 }