Make Iterator implement Enumerable

badger/badger.go

@@ -151,22 +151,12 @@ func (txn *bTxn) prefixIterator(prefix []byte, reverse, keysOnly bool) *prefixIt
 	opt.Prefix = prefix
 	opt.PrefetchValues = !keysOnly
 
-	it := txn.t.NewIterator(opt)
-	if opt.Reverse {
-		it.Seek(bytesPrefixEnd(prefix))
-		if !it.ValidForPrefix(prefix) {
-			it.Next()
-		}
-	} else {
-		it.Seek(prefix) // all badger iterators must start with a rewind
-
-	}
-
 	return &prefixIterator{
-		i:        it,
+		i:        txn.t.NewIterator(opt),
 		prefix:   prefix,
 		reverse:  reverse,
 		keysOnly: keysOnly,
+		reset:    true,
 	}
 }
 
@@ -175,26 +165,13 @@ func (txn *bTxn) rangeIterator(start, end []byte, reverse, keysOnky bool) *range
 	opt.Reverse = reverse
 	opt.PrefetchValues = !keysOnky
 
-	it := txn.t.NewIterator(opt)
-	// all badger iterators must start with a seek/rewind. This is valid
-	// even if start is nil.
-	if !reverse {
-		it.Seek(start)
-	} else {
-		it.Seek(end)
-		// if we seeked to the end and its an exact match to the end marker
-		// go next. This is because ranges are [start, end) (exlusive)
-		if it.Valid() && equal(it.Item().Key(), end) {
-			it.Next()
-		}
-	}
-
 	return &rangeIterator{
-		i:        it,
+		i:        txn.t.NewIterator(opt),
 		start:    start,
 		end:      end,
 		reverse:  reverse,
 		keysOnly: keysOnky,
+		reset:    true,
 	}
 }
 
@@ -229,14 +206,42 @@ type rangeIterator struct {
 	end      []byte
 	reverse  bool
 	keysOnly bool
-	closer   func() error
+	// reset is a mutatuble property that indicates whether the iterator should be
+	// returned to the beginning on the next [Next] call.
+	reset  bool
+	closer func() error
+}
+
+func (it *rangeIterator) Reset() {
+	it.reset = true
+}
+
+// restart returns the iterator back to it's initial location at time of construction,
+// allowing re-iteration of the underlying data.
+func (it *rangeIterator) restart() (bool, error) {
+	it.reset = false
+
+	if it.reverse {
+		hasValue, err := it.Seek(it.end)
+		if !hasValue || err != nil {
+			return false, err
+		}
+
+		if it.valid() && equal(it.i.Item().Key(), it.end) {
+			return it.Next()
+		}
+
+		return true, nil
+	} else {
+		return it.Seek(it.start)
+	}
 }
 
 func (it *rangeIterator) Domain() (start []byte, end []byte) {
 	return it.start, it.end
 }
 
-func (it *rangeIterator) Valid() bool {
+func (it *rangeIterator) valid() bool {
 	if !it.i.Valid() {
 		return false
 	}
@@ -252,8 +257,13 @@ func (it *rangeIterator) Valid() bool {
 	return true
 }
 
-func (it *rangeIterator) Next() {
+func (it *rangeIterator) Next() (bool, error) {
+	if it.reset {
+		return it.restart()
+	}
+
 	it.i.Next()
+	return it.valid(), nil
 }
 
 func (it *rangeIterator) Key() []byte {
@@ -268,8 +278,13 @@ func (it *rangeIterator) Value() ([]byte, error) {
 	return it.i.Item().ValueCopy(nil)
 }
 
-func (it *rangeIterator) Seek(target []byte) {
+func (it *rangeIterator) Seek(target []byte) (bool, error) {
+	// Clear the reset property, else if Next was call following Seek,
+	// Next may incorrectly return to the beginning.
+	it.reset = false
+
 	it.i.Seek(target)
+	return it.valid(), nil
 }
 
 func (it *rangeIterator) Close(ctx context.Context) error {
@@ -289,19 +304,48 @@ type prefixIterator struct {
 	prefix   []byte
 	reverse  bool
 	keysOnly bool
-	closer   func() error
+	// reset is a mutatuble property that indicates whether the iterator should be
+	// returned to the beginning on the next [Next] call.
+	reset  bool
+	closer func() error
+}
+
+func (it *prefixIterator) Reset() {
+	it.reset = true
+}
+
+// restart returns the iterator back to it's initial location at time of construction,
+// allowing re-iteration of the underlying data.
+func (it *prefixIterator) restart() (bool, error) {
+	it.reset = false
+
+	if it.reverse {
+		hasItem, err := it.Seek(bytesPrefixEnd(it.prefix))
+		if !hasItem || err != nil {
+			return false, err
+		}
+
+		if !it.i.ValidForPrefix(it.prefix) {
+			return it.Next()
+		}
+
+		return true, nil
+	} else {
+		return it.Seek(it.prefix)
+	}
 }
 
 func (it *prefixIterator) Domain() (start []byte, end []byte) {
 	return it.prefix, it.prefix
 }
 
-func (it *prefixIterator) Valid() bool {
-	return it.i.ValidForPrefix(it.prefix)
-}
+func (it *prefixIterator) Next() (bool, error) {
+	if it.reset {
+		return it.restart()
+	}
 
-func (it *prefixIterator) Next() {
 	it.i.Next()
+	return it.i.ValidForPrefix(it.prefix), nil
 }
 
 func (it *prefixIterator) Key() []byte {
@@ -315,8 +359,13 @@ func (it *prefixIterator) Value() ([]byte, error) {
 	return it.i.Item().ValueCopy(nil)
 }
 
-func (it *prefixIterator) Seek(target []byte) {
+func (it *prefixIterator) Seek(target []byte) (bool, error) {
+	// Clear the reset property, else if Next was call following Seek,
+	// Next may incorrectly return to the beginning.
+	it.reset = false
+
 	it.i.Seek(target)
+	return it.i.ValidForPrefix(it.prefix), nil
 }
 
 func (it *prefixIterator) Close(ctx context.Context) error {

go.mod

@@ -23,6 +23,7 @@ require (
 	github.com/kr/pretty v0.2.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/sourcenetwork/immutable v0.3.0 // indirect
 	github.com/tidwall/btree v1.7.0 // indirect
 	go.opencensus.io v0.22.5 // indirect
 	golang.org/x/net v0.7.0 // indirect

go.sum

@@ -50,6 +50,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sourcenetwork/immutable v0.3.0 h1:gHPtGvLrTBTK5YpDAhMU+u+S8v1F6iYmc3nbZLryMdc=
+github.com/sourcenetwork/immutable v0.3.0/go.mod h1:GD7ceuh/HD7z6cdIwzKK2ctzgZ1qqYFJpsFp+8qYnbI=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=

kv.go

@@ -89,23 +89,18 @@ type Writer interface {
 
 // Iterator is a read-only iterator that allows iteration over the underlying
 // store (or a part of it).
+//
+// Iterator implements the [Enumerable](https://github.com/sourcenetwork/immutable/blob/main/enumerable/enumerable.go)
+// allowing instances of this type to make use of the various utilities within that package.
 type Iterator interface {
 	// The behaviour of Domain is undefined and varies heavily depending on
 	// which store implementation is being iterated over:
 	// https://github.com/sourcenetwork/corekv/issues/31
 	Domain() (start []byte, end []byte)
 
-	// Valid returns true if the iterator can return a valid item from
-	// `Value`.
-	//
-	// It will return false in all other cases, for example if the store is
-	// empty, or the iterator has passed the end of the range permitted by it's
-	// initialization options.
-	Valid() bool
-
-	// Next moves the iterator forward, whether there is a valid item available or
-	// not.
-	Next()
+	// Next attempts to move the iterator forward, it will return `true` if it was successful,
+	// otherwise `false`.
+	Next() (bool, error)
 
 	// Key returns the key at the current iterator location.
 	//
@@ -119,13 +114,18 @@ type Iterator interface {
 	// https://github.com/sourcenetwork/corekv/issues/37
 	Value() ([]byte, error)
 
-	// Seek moves the iterator to the given key, if and exact match is not found, the
+	// Seek moves the iterator to the given key, if an exact match is not found, the
 	// iterator will progress to the next valid value (depending on the `Reverse` option).
 	//
+	// Seek will return `true` if it found a valid item, otherwise `false`.
+	//
 	// Seek will not seek to values outside of the constraints provided in [IterOptions],
 	// unless using the badger store due to bug:
 	// https://github.com/sourcenetwork/corekv/issues/38
-	Seek([]byte)
+	Seek([]byte) (bool, error)
+
+	// Reset resets the iterator, allowing for re-iteration.
+	Reset()
 
 	// Close releases the iterator.
 	Close(ctx context.Context) error