Merge pull request #564 from strib/tree-walker-use-seen-cache

plumbing: use `seen` map in tree walker

1 files changed, 11 insertions, 5 deletions

diff --git a/plumbing/object/tree.go b/plumbing/object/tree.go
index 512db9f..44ac720 100644
--- a/plumbing/object/tree.go
+++ b/plumbing/object/tree.go
@@ -297,9 +297,10 @@ func (iter *treeEntryIter) Next() (TreeEntry, error) {
 
 // TreeWalker provides a means of walking through all of the entries in a Tree.
 type TreeWalker struct {
-	stack     []treeEntryIter
+	stack     []*treeEntryIter
 	base      string
 	recursive bool
+	seen      map[plumbing.Hash]bool
 
 	s storer.EncodedObjectStorer
 	t *Tree
@@ -309,13 +310,14 @@ type TreeWalker struct {
 //
 // It is the caller's responsibility to call Close() when finished with the
 // tree walker.
-func NewTreeWalker(t *Tree, recursive bool) *TreeWalker {
-	stack := make([]treeEntryIter, 0, startingStackSize)
-	stack = append(stack, treeEntryIter{t, 0})
+func NewTreeWalker(t *Tree, recursive bool, seen map[plumbing.Hash]bool) *TreeWalker {
+	stack := make([]*treeEntryIter, 0, startingStackSize)
+	stack = append(stack, &treeEntryIter{t, 0})
 
 	return &TreeWalker{
 		stack:     stack,
 		recursive: recursive,
+		seen:      seen,
 
 		s: t.s,
 		t: t,
@@ -358,6 +360,10 @@ func (w *TreeWalker) Next() (name string, entry TreeEntry, err error) {
 			return
 		}
 
+		if w.seen[entry.Hash] {
+			continue
+		}
+
 		if entry.Mode == filemode.Dir {
 			obj, err = GetTree(w.s, entry.Hash)
 		}
@@ -377,7 +383,7 @@ func (w *TreeWalker) Next() (name string, entry TreeEntry, err error) {
 	}
 
 	if t, ok := obj.(*Tree); ok {
-		w.stack = append(w.stack, treeEntryIter{t, 0})
+		w.stack = append(w.stack, &treeEntryIter{t, 0})
 		w.base = path.Join(w.base, entry.Name)
 	}