diff --git a/store/cachekvstore.go b/store/cachekvstore.go index e075e7847..6c5cc9542 100644 --- a/store/cachekvstore.go +++ b/store/cachekvstore.go @@ -5,6 +5,7 @@ import ( "sort" "sync" + sdk "github.com/cosmos/cosmos-sdk/types" cmn "github.com/tendermint/tmlibs/common" ) @@ -134,6 +135,16 @@ func (ci *cacheKVStore) ReverseIterator(start, end []byte) Iterator { return ci.iterator(start, end, false) } +// Implements KVStore. +func (ci *cacheKVStore) SubspaceIterator(prefix []byte) Iterator { + return ci.iterator(prefix, sdk.PrefixEndBytes(prefix), true) +} + +// Implements KVStore. +func (ci *cacheKVStore) ReverseSubspaceIterator(prefix []byte) Iterator { + return ci.iterator(prefix, sdk.PrefixEndBytes(prefix), false) +} + func (ci *cacheKVStore) iterator(start, end []byte, ascending bool) Iterator { var parent, cache Iterator if ascending { diff --git a/store/dbstoreadapter.go b/store/dbstoreadapter.go index a1b580fd4..10c175a99 100644 --- a/store/dbstoreadapter.go +++ b/store/dbstoreadapter.go @@ -19,5 +19,13 @@ func (dsa dbStoreAdapter) CacheWrap() CacheWrap { return NewCacheKVStore(dsa) } +func (dsa dbStoreAdapter) SubspaceIterator(prefix []byte) Iterator { + return dsa.Iterator(prefix, sdk.PrefixEndBytes(prefix)) +} + +func (dsa dbStoreAdapter) ReverseSubspaceIterator(prefix []byte) Iterator { + return dsa.ReverseIterator(prefix, sdk.PrefixEndBytes(prefix)) +} + // dbm.DB implements KVStore so we can CacheKVStore it. var _ KVStore = dbStoreAdapter{dbm.DB(nil)} diff --git a/store/iavlstore.go b/store/iavlstore.go index e736fbda2..7d006d3c1 100644 --- a/store/iavlstore.go +++ b/store/iavlstore.go @@ -119,18 +119,21 @@ func (st *iavlStore) Iterator(start, end []byte) Iterator { return newIAVLIterator(st.tree.Tree(), start, end, true) } -func (st *iavlStore) Subspace(prefix []byte) Iterator { - end := make([]byte, len(prefix)) - copy(end, prefix) - end[len(end)-1]++ - return st.Iterator(prefix, end) -} - -// Implements IterKVStore. +// Implements KVStore. func (st *iavlStore) ReverseIterator(start, end []byte) Iterator { return newIAVLIterator(st.tree.Tree(), start, end, false) } +// Implements KVStore. +func (st *iavlStore) SubspaceIterator(prefix []byte) Iterator { + return st.Iterator(prefix, sdk.PrefixEndBytes(prefix)) +} + +// Implements KVStore. +func (st *iavlStore) ReverseSubspaceIterator(prefix []byte) Iterator { + return st.ReverseIterator(prefix, sdk.PrefixEndBytes(prefix)) +} + // Query implements ABCI interface, allows queries // // by default we will return from (latest height -1), @@ -339,6 +342,9 @@ func (iter *iavlIterator) assertIsValid() { //---------------------------------------- func cp(bz []byte) (ret []byte) { + if bz == nil { + return nil + } ret = make([]byte, len(bz)) copy(ret, bz) return ret diff --git a/store/iavlstore_test.go b/store/iavlstore_test.go index 7adae625e..824617c26 100644 --- a/store/iavlstore_test.go +++ b/store/iavlstore_test.go @@ -73,16 +73,74 @@ func TestIAVLIterator(t *testing.T) { iavlStore := newIAVLStore(tree, numHistory) iter := iavlStore.Iterator([]byte("aloha"), []byte("hellz")) expected := []string{"aloha", "hello"} - for i := 0; iter.Valid(); iter.Next() { + var i int + + for i = 0; iter.Valid(); iter.Next() { expectedKey := expected[i] key, value := iter.Key(), iter.Value() assert.EqualValues(t, key, expectedKey) assert.EqualValues(t, value, treeData[expectedKey]) i += 1 } + assert.Equal(t, len(expected), i) + + iter = iavlStore.Iterator([]byte("golang"), []byte("rocks")) + expected = []string{"hello"} + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, treeData[expectedKey]) + i += 1 + } + assert.Equal(t, len(expected), i) + + iter = iavlStore.Iterator(nil, []byte("golang")) + expected = []string{"aloha"} + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, treeData[expectedKey]) + i += 1 + } + assert.Equal(t, len(expected), i) + + iter = iavlStore.Iterator(nil, []byte("shalom")) + expected = []string{"aloha", "hello"} + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, treeData[expectedKey]) + i += 1 + } + assert.Equal(t, len(expected), i) + + iter = iavlStore.Iterator(nil, nil) + expected = []string{"aloha", "hello"} + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, treeData[expectedKey]) + i += 1 + } + assert.Equal(t, len(expected), i) + + iter = iavlStore.Iterator([]byte("golang"), nil) + expected = []string{"hello"} + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, treeData[expectedKey]) + i += 1 + } + assert.Equal(t, len(expected), i) } -func TestIAVLSubspace(t *testing.T) { +func TestIAVLSubspaceIterator(t *testing.T) { db := dbm.NewMemDB() tree, _ := newTree(t, db) iavlStore := newIAVLStore(tree, numHistory) @@ -90,16 +148,114 @@ func TestIAVLSubspace(t *testing.T) { iavlStore.Set([]byte("test1"), []byte("test1")) iavlStore.Set([]byte("test2"), []byte("test2")) iavlStore.Set([]byte("test3"), []byte("test3")) + iavlStore.Set([]byte{byte(55), byte(255), byte(255), byte(0)}, []byte("test4")) + iavlStore.Set([]byte{byte(55), byte(255), byte(255), byte(1)}, []byte("test4")) + iavlStore.Set([]byte{byte(55), byte(255), byte(255), byte(255)}, []byte("test4")) + iavlStore.Set([]byte{byte(255), byte(255), byte(0)}, []byte("test4")) + iavlStore.Set([]byte{byte(255), byte(255), byte(1)}, []byte("test4")) + iavlStore.Set([]byte{byte(255), byte(255), byte(255)}, []byte("test4")) - iter := iavlStore.Subspace([]byte("test")) + i := 0 + + iter := iavlStore.SubspaceIterator([]byte("test")) expected := []string{"test1", "test2", "test3"} - for i := 0; iter.Valid(); iter.Next() { + for i = 0; iter.Valid(); iter.Next() { expectedKey := expected[i] key, value := iter.Key(), iter.Value() assert.EqualValues(t, key, expectedKey) assert.EqualValues(t, value, expectedKey) i += 1 } + assert.Equal(t, len(expected), i) + + iter = iavlStore.SubspaceIterator([]byte{byte(55), byte(255), byte(255)}) + expected2 := [][]byte{ + []byte{byte(55), byte(255), byte(255), byte(0)}, + []byte{byte(55), byte(255), byte(255), byte(1)}, + []byte{byte(55), byte(255), byte(255), byte(255)}, + } + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected2[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, []byte("test4")) + i += 1 + } + assert.Equal(t, len(expected), i) + + iter = iavlStore.SubspaceIterator([]byte{byte(255), byte(255)}) + expected2 = [][]byte{ + []byte{byte(255), byte(255), byte(0)}, + []byte{byte(255), byte(255), byte(1)}, + []byte{byte(255), byte(255), byte(255)}, + } + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected2[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, []byte("test4")) + i += 1 + } + assert.Equal(t, len(expected), i) +} + +func TestIAVLReverseSubspaceIterator(t *testing.T) { + db := dbm.NewMemDB() + tree, _ := newTree(t, db) + iavlStore := newIAVLStore(tree, numHistory) + + iavlStore.Set([]byte("test1"), []byte("test1")) + iavlStore.Set([]byte("test2"), []byte("test2")) + iavlStore.Set([]byte("test3"), []byte("test3")) + iavlStore.Set([]byte{byte(55), byte(255), byte(255), byte(0)}, []byte("test4")) + iavlStore.Set([]byte{byte(55), byte(255), byte(255), byte(1)}, []byte("test4")) + iavlStore.Set([]byte{byte(55), byte(255), byte(255), byte(255)}, []byte("test4")) + iavlStore.Set([]byte{byte(255), byte(255), byte(0)}, []byte("test4")) + iavlStore.Set([]byte{byte(255), byte(255), byte(1)}, []byte("test4")) + iavlStore.Set([]byte{byte(255), byte(255), byte(255)}, []byte("test4")) + + i := 0 + + iter := iavlStore.ReverseSubspaceIterator([]byte("test")) + expected := []string{"test3", "test2", "test1"} + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, expectedKey) + i += 1 + } + assert.Equal(t, len(expected), i) + + iter = iavlStore.ReverseSubspaceIterator([]byte{byte(55), byte(255), byte(255)}) + expected2 := [][]byte{ + []byte{byte(55), byte(255), byte(255), byte(255)}, + []byte{byte(55), byte(255), byte(255), byte(1)}, + []byte{byte(55), byte(255), byte(255), byte(0)}, + } + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected2[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, []byte("test4")) + i += 1 + } + assert.Equal(t, len(expected), i) + + iter = iavlStore.ReverseSubspaceIterator([]byte{byte(255), byte(255)}) + expected2 = [][]byte{ + []byte{byte(255), byte(255), byte(255)}, + []byte{byte(255), byte(255), byte(1)}, + []byte{byte(255), byte(255), byte(0)}, + } + for i = 0; iter.Valid(); iter.Next() { + expectedKey := expected2[i] + key, value := iter.Key(), iter.Value() + assert.EqualValues(t, key, expectedKey) + assert.EqualValues(t, value, []byte("test4")) + i += 1 + } + assert.Equal(t, len(expected), i) } func TestIAVLStoreQuery(t *testing.T) { diff --git a/types/store.go b/types/store.go index ef2e9154b..7b73570ca 100644 --- a/types/store.go +++ b/types/store.go @@ -111,6 +111,14 @@ type KVStore interface { // CONTRACT: No writes may happen within a domain while an iterator exists over it. ReverseIterator(start, end []byte) Iterator + // Iterator over all the keys with a certain prefix in ascending order. + // CONTRACT: No writes may happen within a domain while an iterator exists over it. + SubspaceIterator(prefix []byte) Iterator + + // Iterator over all the keys with a certain prefix in descending order. + // CONTRACT: No writes may happen within a domain while an iterator exists over it. + ReverseSubspaceIterator(prefix []byte) Iterator + // TODO Not yet implemented. // CreateSubKVStore(key *storeKey) (KVStore, error) @@ -222,3 +230,29 @@ func (key *KVStoreKey) Name() string { func (key *KVStoreKey) String() string { return fmt.Sprintf("KVStoreKey{%p, %s}", key, key.name) } + +// PrefixEndBytes returns the []byte that would end a +// range query for all []byte with a certain prefix +// Deals with last byte of prefix being FF without overflowing +func PrefixEndBytes(prefix []byte) []byte { + if prefix == nil { + return nil + } + + end := make([]byte, len(prefix)) + copy(end, prefix) + + for { + if end[len(end)-1] != byte(255) { + end[len(end)-1]++ + break + } else { + end = end[:len(end)-1] + if len(end) == 0 { + end = nil + break + } + } + } + return end +} diff --git a/types/store_test.go b/types/store_test.go new file mode 100644 index 000000000..43dd1f5d3 --- /dev/null +++ b/types/store_test.go @@ -0,0 +1,29 @@ +package types + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestPrefixEndBytes(t *testing.T) { + assert := assert.New(t) + + var testCases = []struct { + prefix []byte + expected []byte + }{ + {[]byte{byte(55), byte(255), byte(255), byte(0)}, []byte{byte(55), byte(255), byte(255), byte(1)}}, + {[]byte{byte(55), byte(255), byte(255), byte(15)}, []byte{byte(55), byte(255), byte(255), byte(16)}}, + {[]byte{byte(55), byte(200), byte(255)}, []byte{byte(55), byte(201)}}, + {[]byte{byte(55), byte(255), byte(255)}, []byte{byte(56)}}, + {[]byte{byte(255), byte(255), byte(255)}, nil}, + {[]byte{byte(255)}, nil}, + {nil, nil}, + } + + for _, test := range testCases { + end := PrefixEndBytes(test.prefix) + assert.Equal(test.expected, end) + } +}