plumbing: storer, new NewMultiReferenceIter

Signed-off-by: Máximo Cuadros <mcuadros@gmail.com>

2 files changed, 89 insertions, 2 deletions

diff --git a/plumbing/storer/reference.go b/plumbing/storer/reference.go
index 5e85a3b..d9fef4a 100644
--- a/plumbing/storer/reference.go
+++ b/plumbing/storer/reference.go
@@ -131,9 +131,27 @@ func (iter *ReferenceSliceIter) Next() (*plumbing.Reference, error) {
 // an error happens or the end of the iter is reached. If ErrStop is sent
 // the iteration is stop but no error is returned. The iterator is closed.
 func (iter *ReferenceSliceIter) ForEach(cb func(*plumbing.Reference) error) error {
+	return forEachReferenceIter(iter, cb)
+}
+
+type bareReferenceIterator interface {
+	Next() (*plumbing.Reference, error)
+	Close()
+}
+
+func forEachReferenceIter(iter bareReferenceIterator, cb func(*plumbing.Reference) error) error {
 	defer iter.Close()
-	for _, r := range iter.series {
-		if err := cb(r); err != nil {
+	for {
+		obj, err := iter.Next()
+		if err != nil {
+			if err == io.EOF {
+				return nil
+			}
+
+			return err
+		}
+
+		if err := cb(obj); err != nil {
 			if err == ErrStop {
 				return nil
 			}
@@ -150,6 +168,52 @@ func (iter *ReferenceSliceIter) Close() {
 	iter.pos = len(iter.series)
 }
 
+// MultiReferenceIter implements ReferenceIter. It iterates over several
+// ReferenceIter,
+//
+// The MultiReferenceIter must be closed with a call to Close() when it is no
+// longer needed.
+type MultiReferenceIter struct {
+	iters []ReferenceIter
+}
+
+// NewMultiReferenceIter returns an reference iterator for the given slice of
+// EncodedObjectIters.
+func NewMultiReferenceIter(iters []ReferenceIter) ReferenceIter {
+	return &MultiReferenceIter{iters: iters}
+}
+
+// Next returns the next reference from the iterator, if one iterator reach
+// io.EOF is removed and the next one is used.
+func (iter *MultiReferenceIter) Next() (*plumbing.Reference, error) {
+	if len(iter.iters) == 0 {
+		return nil, io.EOF
+	}
+
+	obj, err := iter.iters[0].Next()
+	if err == io.EOF {
+		iter.iters[0].Close()
+		iter.iters = iter.iters[1:]
+		return iter.Next()
+	}
+
+	return obj, err
+}
+
+// ForEach call the cb function for each reference contained on this iter until
+// an error happens or the end of the iter is reached. If ErrStop is sent
+// the iteration is stop but no error is returned. The iterator is closed.
+func (iter *MultiReferenceIter) ForEach(cb func(*plumbing.Reference) error) error {
+	return forEachReferenceIter(iter, cb)
+}
+
+// Close releases any resources used by the iterator.
+func (iter *MultiReferenceIter) Close() {
+	for _, i := range iter.iters {
+		i.Close()
+	}
+}
+
 // ResolveReference resolves a SymbolicReference to a HashReference.
 func ResolveReference(s ReferenceStorer, n plumbing.ReferenceName) (*plumbing.Reference, error) {
 	r, err := s.Reference(n)
diff --git a/plumbing/storer/reference_test.go b/plumbing/storer/reference_test.go
index 490ec95..1d02c22 100644
--- a/plumbing/storer/reference_test.go
+++ b/plumbing/storer/reference_test.go
@@ -172,3 +172,26 @@ func (s *ReferenceSuite) TestReferenceFilteredIterForEachStop(c *C) {
 
 	c.Assert(count, Equals, 1)
 }
+
+func (s *ReferenceSuite) TestMultiReferenceIterForEach(c *C) {
+	i := NewMultiReferenceIter(
+		[]ReferenceIter{
+			NewReferenceSliceIter([]*plumbing.Reference{
+				plumbing.NewReferenceFromStrings("foo", "foo"),
+			}),
+			NewReferenceSliceIter([]*plumbing.Reference{
+				plumbing.NewReferenceFromStrings("bar", "bar"),
+			}),
+		},
+	)
+
+	var result []string
+	err := i.ForEach(func(r *plumbing.Reference) error {
+		result = append(result, r.Name().String())
+		return nil
+	})
+
+	c.Assert(err, IsNil)
+	c.Assert(result, HasLen, 2)
+	c.Assert(result, DeepEquals, []string{"foo", "bar"})
+}