storage: filesystem iter implementation

4 files changed, 201 insertions, 40 deletions

diff --git a/fixtures/fixtures.go b/fixtures/fixtures.go
index 7b72650..b85174c 100644
--- a/fixtures/fixtures.go
+++ b/fixtures/fixtures.go
@@ -23,15 +23,18 @@ var fixtures = []*Fixture{{
 	Head:         core.NewHash("6ecf0ef2c2dffb796033e5a02219af86ec6584e5"),
 	PackfileHash: core.NewHash("a3fed42da1e8189a077c0e6846c040dcf73fc9dd"),
 	DotGitHash:   core.NewHash("0a00a25543e6d732dbf4e8e9fec55c8e65fc4e8d"),
+	ObjectsCount: 31,
 }, {
 	Tags:         []string{"ref-delta"},
 	URL:          "https://github.com/git-fixtures/basic",
 	Head:         core.NewHash("6ecf0ef2c2dffb796033e5a02219af86ec6584e5"),
 	PackfileHash: core.NewHash("c544593473465e6315ad4182d04d366c4592b829"),
+	ObjectsCount: 31,
 }, {
-	Tags:       []string{".git", "unpacked", "multi-packfile"},
-	URL:        "https://github.com/src-d/go-git.git",
-	DotGitHash: core.NewHash("174be6bd4292c18160542ae6dc6704b877b8a01a"),
+	Tags:         []string{".git", "unpacked", "multi-packfile"},
+	URL:          "https://github.com/src-d/go-git.git",
+	DotGitHash:   core.NewHash("174be6bd4292c18160542ae6dc6704b877b8a01a"),
+	ObjectsCount: 2133,
 }, {
 	URL:          "https://github.com/spinnaker/spinnaker",
 	Head:         core.NewHash("06ce06d0fc49646c4de733c45b7788aabad98a6f"),
@@ -76,6 +79,7 @@ type Fixture struct {
 	Head         core.Hash
 	PackfileHash core.Hash
 	DotGitHash   core.Hash
+	ObjectsCount int32
 }
 
 func (f *Fixture) Packfile() io.ReadSeeker {
diff --git a/formats/packfile/decoder.go b/formats/packfile/decoder.go
index 6a0aa71..23a8e1a 100644
--- a/formats/packfile/decoder.go
+++ b/formats/packfile/decoder.go
@@ -91,20 +91,17 @@ func (d *Decoder) doDecode() error {
 
 func (d *Decoder) readObjects(count uint32) error {
 	for i := 0; i < int(count); i++ {
-		obj, err := d.readObject()
+		_, err := d.ReadObject()
 		if err != nil {
 			return err
 		}
-
-		if _, err := d.tx.Set(obj); err != nil {
-			return err
-		}
 	}
 
 	return nil
 }
 
-func (d *Decoder) readObject() (core.Object, error) {
+// ReadObject reads a object from the stream and return it
+func (d *Decoder) ReadObject() (core.Object, error) {
 	h, err := d.s.NextObjectHeader()
 	if err != nil {
 		return nil, err
@@ -130,9 +127,31 @@ func (d *Decoder) readObject() (core.Object, error) {
 	}
 
 	d.remember(obj, h.Offset, crc)
+
+	if _, err := d.tx.Set(obj); err != nil {
+		return nil, err
+	}
+
 	return obj, nil
 }
 
+// ReadObjectAt reads an object at the given location
+func (d *Decoder) ReadObjectAt(offset int64) (core.Object, error) {
+	beforeJump, err := d.s.Seek(offset)
+	if err != nil {
+		return nil, err
+	}
+
+	defer func() {
+		_, seekErr := d.s.Seek(beforeJump)
+		if err == nil {
+			err = seekErr
+		}
+	}()
+
+	return d.ReadObject()
+}
+
 func (d *Decoder) fillRegularObjectContent(obj core.Object) (uint32, error) {
 	w, err := obj.Writer()
 	if err != nil {
@@ -195,23 +214,6 @@ func (d *Decoder) recallByHash(h core.Hash) (core.Object, error) {
 	return d.tx.Get(core.AnyObject, h)
 }
 
-// ReadObjectAt reads an object at the given location
-func (d *Decoder) ReadObjectAt(offset int64) (core.Object, error) {
-	beforeJump, err := d.s.Seek(offset)
-	if err != nil {
-		return nil, err
-	}
-
-	defer func() {
-		_, seekErr := d.s.Seek(beforeJump)
-		if err == nil {
-			err = seekErr
-		}
-	}()
-
-	return d.readObject()
-}
-
 // Offsets returns the objects read offset
 func (d *Decoder) Offsets() map[core.Hash]int64 {
 	i := make(map[core.Hash]int64, len(d.offsets))
diff --git a/storage/filesystem/object.go b/storage/filesystem/object.go
index 1da24e5..379c359 100644
--- a/storage/filesystem/object.go
+++ b/storage/filesystem/object.go
@@ -11,6 +11,7 @@ import (
 	"gopkg.in/src-d/go-git.v4/formats/packfile"
 	"gopkg.in/src-d/go-git.v4/storage/filesystem/internal/dotgit"
 	"gopkg.in/src-d/go-git.v4/storage/memory"
+	"gopkg.in/src-d/go-git.v4/utils/fs"
 )
 
 // ObjectStorage is an implementation of core.ObjectStorage that stores
@@ -165,34 +166,50 @@ func (s *ObjectStorage) findObjectInPackfile(h core.Hash) (core.Hash, int64) {
 // Iter returns an iterator for all the objects in the packfile with the
 // given type.
 func (s *ObjectStorage) Iter(t core.ObjectType) (core.ObjectIter, error) {
-	var objects []core.Object
+	objects, err := s.dir.Objects()
+	if err != nil {
+		return nil, err
+	}
+
+	seen := make(map[core.Hash]bool, 0)
+	var iters []core.ObjectIter
+	if len(objects) != 0 {
+		iters = append(iters, &objectsIter{s: s, t: t, h: objects})
+		seen = hashListAsMap(objects)
+	}
 
-	hashes, err := s.dir.Objects()
+	packi, err := s.buildPackfileIters(t, seen)
 	if err != nil {
 		return nil, err
 	}
 
-	for _, hash := range hashes {
-		object, err := s.getFromUnpacked(hash)
+	iters = append(iters, packi...)
+	return core.NewMultiObjectIter(iters), nil
+}
+
+func (s *ObjectStorage) buildPackfileIters(
+	t core.ObjectType, seen map[core.Hash]bool) ([]core.ObjectIter, error) {
+	packs, err := s.dir.ObjectPacks()
+	if err != nil {
+		return nil, err
+	}
+
+	var iters []core.ObjectIter
+	for _, h := range packs {
+		pack, err := s.dir.ObjectPack(h)
 		if err != nil {
 			return nil, err
 		}
-		if object.Type() == t {
-			objects = append(objects, object)
-		}
-	}
 
-	for hash := range s.index {
-		object, err := s.getFromPackfile(hash)
+		iter, err := newPackfileIter(pack, t, seen)
 		if err != nil {
 			return nil, err
 		}
-		if t == core.AnyObject || object.Type() == t {
-			objects = append(objects, object)
-		}
+
+		iters = append(iters, iter)
 	}
 
-	return core.NewObjectSliceIter(objects), nil
+	return iters, nil
 }
 
 func (o *ObjectStorage) Begin() core.TxObjectStorage {
@@ -233,3 +250,103 @@ func (i index) Decode(r io.Reader) error {
 
 	return nil
 }
+
+type packfileIter struct {
+	f fs.File
+	d *packfile.Decoder
+	t core.ObjectType
+
+	seen     map[core.Hash]bool
+	position uint32
+	total    uint32
+}
+
+func newPackfileIter(
+	f fs.File,
+	t core.ObjectType,
+	seen map[core.Hash]bool,
+) (core.ObjectIter, error) {
+	s := packfile.NewScanner(f)
+	_, total, err := s.Header()
+	if err != nil {
+		return nil, err
+	}
+
+	d := packfile.NewDecoder(s, memory.NewStorage().ObjectStorage())
+	return &packfileIter{f: f, d: d, t: t, total: total, seen: seen}, nil
+}
+
+func (iter *packfileIter) Next() (core.Object, error) {
+	if iter.position >= iter.total {
+		return nil, io.EOF
+	}
+
+	obj, err := iter.d.ReadObject()
+	if err != nil {
+		return nil, err
+	}
+
+	iter.position++
+	if iter.seen[obj.Hash()] {
+		return iter.Next()
+	}
+
+	if iter.t != core.AnyObject && iter.t != obj.Type() {
+		return iter.Next()
+	}
+
+	return obj, nil
+}
+
+// ForEach is never called since is used inside of a MultiObjectIterator
+func (iter *packfileIter) ForEach(cb func(core.Object) error) error {
+	return nil
+}
+
+func (iter *packfileIter) Close() {
+	iter.f.Close()
+	iter.d.Close()
+}
+
+type objectsIter struct {
+	s *ObjectStorage
+	t core.ObjectType
+	h []core.Hash
+}
+
+func (iter *objectsIter) Next() (core.Object, error) {
+	if len(iter.h) == 0 {
+		return nil, io.EOF
+	}
+
+	obj, err := iter.s.getFromUnpacked(iter.h[0])
+	iter.h = iter.h[1:]
+
+	if err != nil {
+		return nil, err
+	}
+
+	if iter.t != core.AnyObject && iter.t != obj.Type() {
+		return iter.Next()
+	}
+
+	return obj, err
+}
+
+// ForEach is never called since is used inside of a MultiObjectIterator
+func (iter *objectsIter) ForEach(cb func(core.Object) error) error {
+	return nil
+}
+
+func (iter *objectsIter) Close() {
+	iter.h = []core.Hash{}
+}
+
+func hashListAsMap(l []core.Hash) map[core.Hash]bool {
+	m := make(map[core.Hash]bool, len(l))
+	for _, h := range l {
+		m[h] = true
+	}
+
+	return m
+}
diff --git a/storage/filesystem/object_test.go b/storage/filesystem/object_test.go
index 99266ab..14c77e4 100644
--- a/storage/filesystem/object_test.go
+++ b/storage/filesystem/object_test.go
@@ -53,3 +53,41 @@ func (s *FsSuite) TestGetFromPackfileMultiplePackfiles(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(obj.Hash(), Equals, expected)
 }
+
+func (s *FsSuite) TestIter(c *C) {
+	fixtures.ByTag(".git").Test(c, func(f *fixtures.Fixture) {
+		fs := f.DotGit()
+		o, err := newObjectStorage(dotgit.New(fs))
+		c.Assert(err, IsNil)
+
+		iter, err := o.Iter(core.AnyObject)
+		c.Assert(err, IsNil)
+
+		var count int32
+		err = iter.ForEach(func(o core.Object) error {
+			count++
+			return nil
+		})
+
+		c.Assert(err, IsNil)
+		c.Assert(count, Equals, f.ObjectsCount)
+	})
+}
+
+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)
+
+		iter, err := o.Iter(core.CommitObject)
+		c.Assert(err, IsNil)
+
+		err = iter.ForEach(func(o core.Object) error {
+			c.Assert(o.Type(), Equals, core.CommitObject)
+			return nil
+		})
+
+		c.Assert(err, IsNil)
+	})
+}