packfile/decoder: speed up packfile iterator when specific type (#200)

2 files changed, 70 insertions, 27 deletions

diff --git a/storage/filesystem/object.go b/storage/filesystem/object.go
index e9b5bb7..1001032 100644
--- a/storage/filesystem/object.go
+++ b/storage/filesystem/object.go
@@ -271,6 +271,10 @@ type packfileIter struct {
 	total    uint32
 }
 
+func NewPackfileIter(f billy.File, t plumbing.ObjectType) (storer.EncodedObjectIter, error) {
+	return newPackfileIter(f, t, make(map[plumbing.Hash]bool))
+}
+
 func newPackfileIter(f billy.File, t plumbing.ObjectType, seen map[plumbing.Hash]bool) (storer.EncodedObjectIter, error) {
 	s := packfile.NewScanner(f)
 	_, total, err := s.Header()
@@ -278,7 +282,7 @@ func newPackfileIter(f billy.File, t plumbing.ObjectType, seen map[plumbing.Hash
 		return nil, err
 	}
 
-	d, err := packfile.NewDecoder(s, memory.NewStorage())
+	d, err := packfile.NewDecoderForType(s, memory.NewStorage(), t)
 	if err != nil {
 		return nil, err
 	}
@@ -294,25 +298,27 @@ func newPackfileIter(f billy.File, t plumbing.ObjectType, seen map[plumbing.Hash
 }
 
 func (iter *packfileIter) Next() (plumbing.EncodedObject, error) {
-	if iter.position >= iter.total {
-		return nil, io.EOF
-	}
+	for {
+		if iter.position >= iter.total {
+			return nil, io.EOF
+		}
 
-	obj, err := iter.d.DecodeObject()
-	if err != nil {
-		return nil, err
-	}
+		obj, err := iter.d.DecodeObject()
+		if err != nil {
+			return nil, err
+		}
 
-	iter.position++
-	if iter.seen[obj.Hash()] {
-		return iter.Next()
-	}
+		iter.position++
+		if obj == nil {
+			continue
+		}
 
-	if iter.t != plumbing.AnyObject && iter.t != obj.Type() {
-		return iter.Next()
-	}
+		if iter.seen[obj.Hash()] {
+			return iter.Next()
+		}
 
-	return obj, nil
+		return obj, nil
+	}
 }
 
 // ForEach is never called since is used inside of a MultiObjectIterator
diff --git a/storage/filesystem/object_test.go b/storage/filesystem/object_test.go
index a3a5e68..6841ac4 100644
--- a/storage/filesystem/object_test.go
+++ b/storage/filesystem/object_test.go
@@ -10,9 +10,17 @@ import (
 
 type FsSuite struct {
 	fixtures.Suite
+	Types []plumbing.ObjectType
 }
 
-var _ = Suite(&FsSuite{})
+var _ = Suite(&FsSuite{
+	Types: []plumbing.ObjectType{
+		plumbing.CommitObject,
+		plumbing.TagObject,
+		plumbing.TreeObject,
+		plumbing.BlobObject,
+	},
+})
 
 func (s *FsSuite) TestGetFromObjectFile(c *C) {
 	fs := fixtures.ByTag(".git").ByTag("unpacked").One().DotGit()
@@ -76,18 +84,47 @@ func (s *FsSuite) TestIter(c *C) {
 
 func (s *FsSuite) TestIterWithType(c *C) {
 	fixtures.ByTag(".git").Test(c, func(f *fixtures.Fixture) {
-		fs := f.DotGit()
-		o, err := newObjectStorage(dotgit.New(fs))
-		c.Assert(err, IsNil)
+		for _, t := range s.Types {
+			fs := f.DotGit()
+			o, err := newObjectStorage(dotgit.New(fs))
+			c.Assert(err, IsNil)
 
-		iter, err := o.IterEncodedObjects(plumbing.CommitObject)
-		c.Assert(err, IsNil)
+			iter, err := o.IterEncodedObjects(t)
+			c.Assert(err, IsNil)
 
-		err = iter.ForEach(func(o plumbing.EncodedObject) error {
-			c.Assert(o.Type(), Equals, plumbing.CommitObject)
-			return nil
-		})
+			err = iter.ForEach(func(o plumbing.EncodedObject) error {
+				c.Assert(o.Type(), Equals, t)
+				return nil
+			})
+
+			c.Assert(err, IsNil)
+		}
 
-		c.Assert(err, IsNil)
 	})
 }
+
+func (s *FsSuite) TestPackfileIter(c *C) {
+	fixtures.ByTag(".git").Test(c, func(f *fixtures.Fixture) {
+		fs := f.DotGit()
+		dg := dotgit.New(fs)
+
+		for _, t := range s.Types {
+			ph, err := dg.ObjectPacks()
+			c.Assert(err, IsNil)
+
+			for _, h := range ph {
+				f, err := dg.ObjectPack(h)
+				c.Assert(err, IsNil)
+				iter, err := NewPackfileIter(f, t)
+				c.Assert(err, IsNil)
+				err = iter.ForEach(func(o plumbing.EncodedObject) error {
+					c.Assert(o.Type(), Equals, t)
+					return nil
+				})
+
+				c.Assert(err, IsNil)
+			}
+		}
+	})
+
+}