148 lines
2.8 KiB
Go
148 lines
2.8 KiB
Go
|
package rocksdb
|
||
|
|
||
|
import (
|
||
|
"bytes"
|
||
|
|
||
|
dbm "github.com/cosmos/cosmos-sdk/db"
|
||
|
"github.com/tecbot/gorocksdb"
|
||
|
)
|
||
|
|
||
|
type rocksDBIterator struct {
|
||
|
source *gorocksdb.Iterator
|
||
|
start, end []byte
|
||
|
isReverse bool
|
||
|
isInvalid bool
|
||
|
// Whether iterator has been advanced to the first element (is fully initialized)
|
||
|
primed bool
|
||
|
}
|
||
|
|
||
|
var _ dbm.Iterator = (*rocksDBIterator)(nil)
|
||
|
|
||
|
func newRocksDBIterator(source *gorocksdb.Iterator, start, end []byte, isReverse bool) *rocksDBIterator {
|
||
|
if isReverse {
|
||
|
if end == nil {
|
||
|
source.SeekToLast()
|
||
|
} else {
|
||
|
source.Seek(end)
|
||
|
if source.Valid() {
|
||
|
eoakey := moveSliceToBytes(source.Key()) // end or after key
|
||
|
if bytes.Compare(end, eoakey) <= 0 {
|
||
|
source.Prev()
|
||
|
}
|
||
|
} else {
|
||
|
source.SeekToLast()
|
||
|
}
|
||
|
}
|
||
|
} else {
|
||
|
if start == nil {
|
||
|
source.SeekToFirst()
|
||
|
} else {
|
||
|
source.Seek(start)
|
||
|
}
|
||
|
}
|
||
|
return &rocksDBIterator{
|
||
|
source: source,
|
||
|
start: start,
|
||
|
end: end,
|
||
|
isReverse: isReverse,
|
||
|
isInvalid: false,
|
||
|
primed: false,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Domain implements Iterator.
|
||
|
func (itr *rocksDBIterator) Domain() ([]byte, []byte) {
|
||
|
return itr.start, itr.end
|
||
|
}
|
||
|
|
||
|
// Valid implements Iterator.
|
||
|
func (itr *rocksDBIterator) Valid() bool {
|
||
|
if !itr.primed {
|
||
|
return false
|
||
|
}
|
||
|
|
||
|
if itr.isInvalid {
|
||
|
return false
|
||
|
}
|
||
|
|
||
|
if !itr.source.Valid() {
|
||
|
itr.isInvalid = true
|
||
|
return false
|
||
|
}
|
||
|
|
||
|
var (
|
||
|
start = itr.start
|
||
|
end = itr.end
|
||
|
key = moveSliceToBytes(itr.source.Key())
|
||
|
)
|
||
|
// If key is end or past it, invalid.
|
||
|
if itr.isReverse {
|
||
|
if start != nil && bytes.Compare(key, start) < 0 {
|
||
|
itr.isInvalid = true
|
||
|
return false
|
||
|
}
|
||
|
} else {
|
||
|
if end != nil && bytes.Compare(key, end) >= 0 {
|
||
|
itr.isInvalid = true
|
||
|
return false
|
||
|
}
|
||
|
}
|
||
|
return true
|
||
|
}
|
||
|
|
||
|
// Key implements Iterator.
|
||
|
func (itr *rocksDBIterator) Key() []byte {
|
||
|
itr.assertIsValid()
|
||
|
return moveSliceToBytes(itr.source.Key())
|
||
|
}
|
||
|
|
||
|
// Value implements Iterator.
|
||
|
func (itr *rocksDBIterator) Value() []byte {
|
||
|
itr.assertIsValid()
|
||
|
return moveSliceToBytes(itr.source.Value())
|
||
|
}
|
||
|
|
||
|
// Next implements Iterator.
|
||
|
func (itr *rocksDBIterator) Next() bool {
|
||
|
if !itr.primed {
|
||
|
itr.primed = true
|
||
|
} else {
|
||
|
if itr.isReverse {
|
||
|
itr.source.Prev()
|
||
|
} else {
|
||
|
itr.source.Next()
|
||
|
}
|
||
|
}
|
||
|
return itr.Valid()
|
||
|
}
|
||
|
|
||
|
// Error implements Iterator.
|
||
|
func (itr *rocksDBIterator) Error() error {
|
||
|
return itr.source.Err()
|
||
|
}
|
||
|
|
||
|
// Close implements Iterator.
|
||
|
func (itr *rocksDBIterator) Close() error {
|
||
|
itr.source.Close()
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
func (itr *rocksDBIterator) assertIsValid() {
|
||
|
if !itr.Valid() {
|
||
|
panic("iterator is invalid")
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// moveSliceToBytes will free the slice and copy out a go []byte
|
||
|
// This function can be applied on *Slice returned from Key() and Value()
|
||
|
// of an Iterator, because they are marked as freed.
|
||
|
func moveSliceToBytes(s *gorocksdb.Slice) []byte {
|
||
|
defer s.Free()
|
||
|
if !s.Exists() {
|
||
|
return nil
|
||
|
}
|
||
|
v := make([]byte, s.Size())
|
||
|
copy(v, s.Data())
|
||
|
return v
|
||
|
}
|