prefix store support object store (#236)

Author: yihuang

store/cachekv/store.go

@@ -89,17 +89,10 @@ func (store *GStore[V]) Get(key []byte) (value V) {
 	return value
 }
 
-func (store *GStore[V]) assertValidValue(value V) {
-	if store.isZero(value) {
-		panic("value is nil")
-	}
-	types.AssertValidValueLength(store.valueLen(value))
-}
-
 // Set implements types.KVStore.
 func (store *GStore[V]) Set(key []byte, value V) {
 	types.AssertValidKey(key)
-	store.assertValidValue(value)
+	types.AssertValidValueGeneric(value, store.isZero, store.valueLen)
 
 	store.mtx.Lock()
 	defer store.mtx.Unlock()

store/gaskv/store.go

@@ -71,17 +71,10 @@ func (gs *GStore[V]) Get(key []byte) (value V) {
 	return value
 }
 
-func (gs *GStore[V]) assertValidValue(value V) {
-	if gs.isZero(value) {
-		panic("value is nil")
-	}
-	types.AssertValidValueLength(gs.valueLen(value))
-}
-
 // Implements KVStore.
 func (gs *GStore[V]) Set(key []byte, value V) {
 	types.AssertValidKey(key)
-	gs.assertValidValue(value)
+	types.AssertValidValueGeneric(value, gs.isZero, gs.valueLen)
 	gs.gasMeter.ConsumeGas(gs.gasConfig.WriteCostFlat, types.GasWriteCostFlatDesc)
 	// TODO overflow-safe math?
 	gs.gasMeter.ConsumeGas(gs.gasConfig.WriteCostPerByte*types.Gas(len(key)), types.GasWritePerByteDesc)

store/prefix/store.go

@@ -8,20 +8,53 @@ import (
 	"cosmossdk.io/store/types"
 )
 
-var _ types.KVStore = Store{}
+type (
+	Store    = GStore[[]byte]
+	ObjStore = GStore[any]
+)
+
+var (
+	_ types.KVStore    = Store{}
+	_ types.ObjKVStore = ObjStore{}
+)
+
+func NewStore(parent types.KVStore, prefix []byte) Store {
+	return NewGStore(
+		parent, prefix,
+		func(v []byte) bool { return v == nil },
+		func(v []byte) int { return len(v) },
+	)
+}
 
-// Store is similar with cometbft/cometbft/libs/db/prefix_db
+func NewObjStore(parent types.ObjKVStore, prefix []byte) ObjStore {
+	return NewGStore(
+		parent, prefix,
+		func(v any) bool { return v == nil },
+		func(v any) int { return 1 },
+	)
+}
+
+// GStore is similar with cometbft/cometbft/libs/db/prefix_db
 // both gives access only to the limited subset of the store
 // for convinience or safety
-type Store struct {
-	parent types.KVStore
+type GStore[V any] struct {
+	parent types.GKVStore[V]
 	prefix []byte
+
+	isZero   func(V) bool
+	valueLen func(V) int
 }
 
-func NewStore(parent types.KVStore, prefix []byte) Store {
-	return Store{
+func NewGStore[V any](
+	parent types.GKVStore[V], prefix []byte,
+	isZero func(V) bool, valueLen func(V) int,
+) GStore[V] {
+	return GStore[V]{
 		parent: parent,
 		prefix: prefix,
+
+		isZero:   isZero,
+		valueLen: valueLen,
 	}
 }
 
@@ -32,7 +65,7 @@ func cloneAppend(bz, tail []byte) (res []byte) {
 	return
 }
 
-func (s Store) key(key []byte) (res []byte) {
+func (s GStore[V]) key(key []byte) (res []byte) {
 	if key == nil {
 		panic("nil key on Store")
 	}
@@ -41,41 +74,41 @@ func (s Store) key(key []byte) (res []byte) {
 }
 
 // Implements Store
-func (s Store) GetStoreType() types.StoreType {
+func (s GStore[V]) GetStoreType() types.StoreType {
 	return s.parent.GetStoreType()
 }
 
 // Implements CacheWrap
-func (s Store) CacheWrap() types.CacheWrap {
-	return cachekv.NewStore(s)
+func (s GStore[V]) CacheWrap() types.CacheWrap {
+	return cachekv.NewGStore(s, s.isZero, s.valueLen)
 }
 
 // Implements KVStore
-func (s Store) Get(key []byte) []byte {
+func (s GStore[V]) Get(key []byte) V {
 	res := s.parent.Get(s.key(key))
 	return res
 }
 
 // Implements KVStore
-func (s Store) Has(key []byte) bool {
+func (s GStore[V]) Has(key []byte) bool {
 	return s.parent.Has(s.key(key))
 }
 
 // Implements KVStore
-func (s Store) Set(key, value []byte) {
+func (s GStore[V]) Set(key []byte, value V) {
 	types.AssertValidKey(key)
-	types.AssertValidValue(value)
+	types.AssertValidValueGeneric(value, s.isZero, s.valueLen)
 	s.parent.Set(s.key(key), value)
 }
 
 // Implements KVStore
-func (s Store) Delete(key []byte) {
+func (s GStore[V]) Delete(key []byte) {
 	s.parent.Delete(s.key(key))
 }
 
 // Implements KVStore
 // Check https://github.com/cometbft/cometbft/blob/master/libs/db/prefix_db.go#L106
-func (s Store) Iterator(start, end []byte) types.Iterator {
+func (s GStore[V]) Iterator(start, end []byte) types.GIterator[V] {
 	newstart := cloneAppend(s.prefix, start)
 
 	var newend []byte
@@ -92,7 +125,7 @@ func (s Store) Iterator(start, end []byte) types.Iterator {
 
 // ReverseIterator implements KVStore
 // Check https://github.com/cometbft/cometbft/blob/master/libs/db/prefix_db.go#L129
-func (s Store) ReverseIterator(start, end []byte) types.Iterator {
+func (s GStore[V]) ReverseIterator(start, end []byte) types.GIterator[V] {
 	newstart := cloneAppend(s.prefix, start)
 
 	var newend []byte
@@ -107,18 +140,18 @@ func (s Store) ReverseIterator(start, end []byte) types.Iterator {
 	return newPrefixIterator(s.prefix, start, end, iter)
 }
 
-var _ types.Iterator = (*prefixIterator)(nil)
+var _ types.Iterator = (*prefixIterator[[]byte])(nil)
 
-type prefixIterator struct {
+type prefixIterator[V any] struct {
 	prefix []byte
 	start  []byte
 	end    []byte
-	iter   types.Iterator
+	iter   types.GIterator[V]
 	valid  bool
 }
 
-func newPrefixIterator(prefix, start, end []byte, parent types.Iterator) *prefixIterator {
-	return &prefixIterator{
+func newPrefixIterator[V any](prefix, start, end []byte, parent types.GIterator[V]) *prefixIterator[V] {
+	return &prefixIterator[V]{
 		prefix: prefix,
 		start:  start,
 		end:    end,
@@ -128,17 +161,17 @@ func newPrefixIterator(prefix, start, end []byte, parent types.Iterator) *prefix
 }
 
 // Implements Iterator
-func (pi *prefixIterator) Domain() ([]byte, []byte) {
+func (pi *prefixIterator[V]) Domain() ([]byte, []byte) {
 	return pi.start, pi.end
 }
 
 // Implements Iterator
-func (pi *prefixIterator) Valid() bool {
+func (pi *prefixIterator[V]) Valid() bool {
 	return pi.valid && pi.iter.Valid()
 }
 
 // Implements Iterator
-func (pi *prefixIterator) Next() {
+func (pi *prefixIterator[V]) Next() {
 	if !pi.valid {
 		panic("prefixIterator invalid, cannot call Next()")
 	}
@@ -150,7 +183,7 @@ func (pi *prefixIterator) Next() {
 }
 
 // Implements Iterator
-func (pi *prefixIterator) Key() (key []byte) {
+func (pi *prefixIterator[V]) Key() (key []byte) {
 	if !pi.valid {
 		panic("prefixIterator invalid, cannot call Key()")
 	}
@@ -162,7 +195,7 @@ func (pi *prefixIterator) Key() (key []byte) {
 }
 
 // Implements Iterator
-func (pi *prefixIterator) Value() []byte {
+func (pi *prefixIterator[V]) Value() V {
 	if !pi.valid {
 		panic("prefixIterator invalid, cannot call Value()")
 	}
@@ -171,13 +204,13 @@ func (pi *prefixIterator) Value() []byte {
 }
 
 // Implements Iterator
-func (pi *prefixIterator) Close() error {
+func (pi *prefixIterator[V]) Close() error {
 	return pi.iter.Close()
 }
 
 // Error returns an error if the prefixIterator is invalid defined by the Valid
 // method.
-func (pi *prefixIterator) Error() error {
+func (pi *prefixIterator[V]) Error() error {
 	if !pi.Valid() {
 		return errors.New("invalid prefixIterator")
 	}

store/types/store.go

@@ -583,3 +583,17 @@ func NewMemoryStoreKeys(names ...string) map[string]*MemoryStoreKey {
 
 	return keys
 }
+
+// NewObjectStoreKeys constructs a new map matching store key names to their
+// respective ObjectStoreKey references.
+// The function will panic if there is a potential conflict in names (see `assertNoPrefix`
+// function for more details).
+func NewObjectStoreKeys(names ...string) map[string]*ObjectStoreKey {
+	assertNoCommonPrefix(names)
+	keys := make(map[string]*ObjectStoreKey)
+	for _, n := range names {
+		keys[n] = NewObjectStoreKey(n)
+	}
+
+	return keys
+}

store/types/validity.go

@@ -27,6 +27,14 @@ func AssertValidValue(value []byte) {
 	AssertValidValueLength(len(value))
 }
 
+// AssertValidValueGeneric checks if the value is valid(value is not nil and within length limit)
+func AssertValidValueGeneric[V any](value V, isZero func(V) bool, valueLen func(V) int) {
+	if isZero(value) {
+		panic("value is nil")
+	}
+	AssertValidValueLength(valueLen(value))
+}
+
 // AssertValidValueLength checks if the value length is within length limit
 func AssertValidValueLength(l int) {
 	if l > MaxValueLength {

types/context.go

@@ -350,7 +350,7 @@ func (c Context) TransientStore(key storetypes.StoreKey) storetypes.KVStore {
 }
 
 // ObjectStore fetches an object store from the MultiStore,
-func (c Context) OjectStore(key storetypes.StoreKey) storetypes.ObjKVStore {
+func (c Context) ObjectStore(key storetypes.StoreKey) storetypes.ObjKVStore {
 	return gaskv.NewObjStore(c.ms.GetObjKVStore(key), c.gasMeter, c.transientKVGasConfig)
 }