core: *Iter.ForEach method

4 files changed, 123 insertions, 1 deletions

diff --git a/core/object.go b/core/object.go
index 36b669e..d777f4a 100644
--- a/core/object.go
+++ b/core/object.go
@@ -139,6 +139,28 @@ func (iter *ObjectLookupIter) Next() (Object, error) {
 	return obj, err
 }
 
+// ForEach call the cb function for each object contained on this iter until
+// an error happends or the end of the iter is reached. If ErrStop is sent
+// the iteration is stop but no error is returned
+func (iter *ObjectLookupIter) ForEach(cb func(Object) error) error {
+	for _, hash := range iter.series {
+		obj, err := iter.storage.Get(hash)
+		if err != nil {
+			return err
+		}
+
+		if err := cb(obj); err != nil {
+			if err == ErrStop {
+				return nil
+			}
+
+			return nil
+		}
+	}
+
+	return nil
+}
+
 // Close releases any resources used by the iterator.
 func (iter *ObjectLookupIter) Close() {
 	iter.pos = len(iter.series)
@@ -173,6 +195,23 @@ func (iter *ObjectSliceIter) Next() (Object, error) {
 	return obj, nil
 }
 
+// ForEach call the cb function for each object contained on this iter until
+// an error happends or the end of the iter is reached. If ErrStop is sent
+// the iteration is stop but no error is returned
+func (iter *ObjectSliceIter) ForEach(cb func(Object) error) error {
+	for _, o := range iter.series {
+		if err := cb(o); err != nil {
+			if err == ErrStop {
+				return nil
+			}
+
+			return nil
+		}
+	}
+
+	return nil
+}
+
 // Close releases any resources used by the iterator.
 func (iter *ObjectSliceIter) Close() {
 	iter.pos = len(iter.series)
diff --git a/core/reference.go b/core/reference.go
index 89852da..bde3ff4 100644
--- a/core/reference.go
+++ b/core/reference.go
@@ -178,6 +178,23 @@ func (iter *ReferenceSliceIter) Next() (*Reference, error) {
 	return obj, nil
 }
 
+// ForEach call the cb function for each reference contained on this iter until
+// an error happends or the end of the iter is reached. If ErrStop is sent
+// the iteration is stop but no error is returned
+func (iter *ReferenceSliceIter) ForEach(cb func(*Reference) error) error {
+	for _, r := range iter.series {
+		if err := cb(r); err != nil {
+			if err == ErrStop {
+				return nil
+			}
+
+			return nil
+		}
+	}
+
+	return nil
+}
+
 // Close releases any resources used by the iterator.
 func (iter *ReferenceSliceIter) Close() {
 	iter.pos = len(iter.series)
diff --git a/core/reference_test.go b/core/reference_test.go
index 5d88f0e..3739a15 100644
--- a/core/reference_test.go
+++ b/core/reference_test.go
@@ -1,6 +1,10 @@
 package core
 
-import . "gopkg.in/check.v1"
+import (
+	"io"
+
+	. "gopkg.in/check.v1"
+)
 
 type ReferenceSuite struct{}
 
@@ -62,3 +66,58 @@ func (s *ReferenceSuite) TestIsTag(c *C) {
 	r := NewHashReference(ReferenceName("refs/tags/v3.1."), ZeroHash)
 	c.Assert(r.IsTag(), Equals, true)
 }
+
+func (s *ReferenceSuite) TestReferenceSliceIterNext(c *C) {
+	slice := []*Reference{
+		NewReferenceFromStrings("foo", "foo"),
+		NewReferenceFromStrings("bar", "bar"),
+	}
+
+	i := NewReferenceSliceIter(slice)
+	foo, err := i.Next()
+	c.Assert(err, IsNil)
+	c.Assert(foo == slice[0], Equals, true)
+
+	bar, err := i.Next()
+	c.Assert(err, IsNil)
+	c.Assert(bar == slice[1], Equals, true)
+
+	empty, err := i.Next()
+	c.Assert(err, Equals, io.EOF)
+	c.Assert(empty, IsNil)
+}
+
+func (s *ReferenceSuite) TestReferenceSliceIterForEach(c *C) {
+	slice := []*Reference{
+		NewReferenceFromStrings("foo", "foo"),
+		NewReferenceFromStrings("bar", "bar"),
+	}
+
+	i := NewReferenceSliceIter(slice)
+	var count int
+	i.ForEach(func(r *Reference) error {
+		c.Assert(r == slice[count], Equals, true)
+		count++
+		return nil
+	})
+
+	c.Assert(count, Equals, 2)
+}
+
+func (s *ReferenceSuite) TestReferenceSliceIterForEachStop(c *C) {
+	slice := []*Reference{
+		NewReferenceFromStrings("foo", "foo"),
+		NewReferenceFromStrings("bar", "bar"),
+	}
+
+	i := NewReferenceSliceIter(slice)
+
+	var count int
+	i.ForEach(func(r *Reference) error {
+		c.Assert(r == slice[count], Equals, true)
+		count++
+		return ErrStop
+	})
+
+	c.Assert(count, Equals, 1)
+}
diff --git a/core/storage.go b/core/storage.go
index cb8f45d..6403f5f 100644
--- a/core/storage.go
+++ b/core/storage.go
@@ -1,5 +1,10 @@
 package core
 
+import "errors"
+
+//ErrStop is used to stop a ForEach function in an Iter
+var ErrStop = errors.New("stop iter")
+
 // Storage storage of objects and references
 type Storage interface {
 	ObjectStorage() ObjectStorage
@@ -17,6 +22,7 @@ type ObjectStorage interface {
 // ObjectIter is a generic closable interface for iterating over objects.
 type ObjectIter interface {
 	Next() (Object, error)
+	ForEach(func(Object) error) error
 	Close()
 }
 
@@ -30,5 +36,6 @@ type ReferenceStorage interface {
 // ReferenceIter is a generic closable interface for iterating over references
 type ReferenceIter interface {
 	Next() (*Reference, error)
+	ForEach(func(*Reference) error) error
 	Close()
 }