Merge pull request cosmos#767 from cosmos/sunny/IAVLsubspace

KVStore Subspaces

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 {

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)}

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

store/iavlstore_test.go

@@ -73,33 +73,189 @@ 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)
 
 	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.Subspace([]byte("test"))
+	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) {

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
+}

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)
+	}
+}