feat(collections): pair keys (#14310)

Co-authored-by: testinginprod <testinginprod@somewhere.idk>
Co-authored-by: Aaron Craelius <aaron@regen.network>

collections/CHANGELOG.md

@@ -34,4 +34,5 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#14134](https://github.com/cosmos/cosmos-sdk/pull/14134) Initialise core (Prefix, KeyEncoder, ValueEncoder, Map).
 * [#14351](https://github.com/cosmos/cosmos-sdk/pull/14351) Add keyset
 * [#14364](https://github.com/cosmos/cosmos-sdk/pull/14364) Add sequence
-* [#14468](https://github.com/cosmos/cosmos-sdk/pull/14468) Add Map.IterateRaw API.
+* [#14468](https://github.com/cosmos/cosmos-sdk/pull/14468) Add Map.IterateRaw API.
+* [#14310](https://github.com/cosmos/cosmos-sdk/pull/14310) Add Pair keys 

collections/collections.go

@@ -89,6 +89,25 @@ type KeyCodec[T any] interface {
 	Stringify(key T) string
 	// KeyType returns a string identifier for the type of the key.
 	KeyType() string
+
+	// MULTIPART keys
+
+	// EncodeNonTerminal writes the key bytes into the buffer.
+	// EncodeNonTerminal is used in multipart keys like Pair
+	// when the part of the key being encoded is not the last one,
+	// and there needs to be a way to distinguish after how many bytes
+	// the first part of the key is finished. The buffer is expected to be
+	// at least as big as SizeNonTerminal(key) returns. It returns
+	// the amount of bytes written.
+	EncodeNonTerminal(buffer []byte, key T) (int, error)
+	// DecodeNonTerminal reads the buffer provided and returns
+	// the key T. DecodeNonTerminal is used in multipart keys
+	// like Pair when the part of the key being decoded is not the
+	// last one. It returns the amount of bytes read.
+	DecodeNonTerminal(buffer []byte) (int, T, error)
+	// SizeNonTerminal returns the maximum size of the key K when used in
+	// multipart keys like Pair.
+	SizeNonTerminal(key T) int
 }
 
 // ValueCodec defines a generic interface which is implemented

collections/collections_test.go

@@ -73,15 +73,25 @@ func deps() (store.KVStoreService, context.Context) {
 }
 
 // checkKeyCodec asserts the correct behaviour of a KeyCodec over the type T.
-func checkKeyCodec[T any](t *testing.T, encoder KeyCodec[T], key T) {
-	buffer := make([]byte, encoder.Size(key))
-	written, err := encoder.Encode(buffer, key)
+func checkKeyCodec[T any](t *testing.T, keyCodec KeyCodec[T], key T) {
+	buffer := make([]byte, keyCodec.Size(key))
+	written, err := keyCodec.Encode(buffer, key)
 	require.NoError(t, err)
 	require.Equal(t, len(buffer), written)
-	read, decodedKey, err := encoder.Decode(buffer)
+	read, decodedKey, err := keyCodec.Decode(buffer)
 	require.NoError(t, err)
 	require.Equal(t, len(buffer), read, "encoded key and read bytes must have same size")
 	require.Equal(t, key, decodedKey, "encoding and decoding produces different keys")
+	// test if terminality is correctly applied
+	pairCodec := PairKeyCodec(keyCodec, StringKey)
+	pairKey := Join(key, "TEST")
+	buffer = make([]byte, pairCodec.Size(pairKey))
+	written, err = pairCodec.Encode(buffer, pairKey)
+	require.NoError(t, err)
+	read, decodedPairKey, err := pairCodec.Decode(buffer)
+	require.NoError(t, err)
+	require.Equal(t, len(buffer), read, "encoded non terminal key and pair key read bytes must have same size")
+	require.Equal(t, pairKey, decodedPairKey, "encoding and decoding produces different keys with non terminal encoding")
 }
 
 // checkValueCodec asserts the correct behaviour of a ValueCodec over the type T.

collections/item.go

@@ -55,8 +55,11 @@ func (i Item[V]) Remove(ctx context.Context) error {
 // noKey defines a KeyCodec which decodes nothing.
 type noKey struct{}
 
-func (noKey) Stringify(_ noKey) string              { return "no_key" }
-func (noKey) KeyType() string                       { return "no_key" }
-func (noKey) Size(_ noKey) int                      { return 0 }
-func (noKey) Encode(_ []byte, _ noKey) (int, error) { return 0, nil }
-func (noKey) Decode(_ []byte) (int, noKey, error)   { return 0, noKey{}, nil }
+func (noKey) Stringify(_ noKey) string                           { return "no_key" }
+func (noKey) KeyType() string                                    { return "no_key" }
+func (noKey) Size(_ noKey) int                                   { return 0 }
+func (noKey) Encode(_ []byte, _ noKey) (int, error)              { return 0, nil }
+func (noKey) Decode(_ []byte) (int, noKey, error)                { return 0, noKey{}, nil }
+func (k noKey) EncodeNonTerminal(_ []byte, _ noKey) (int, error) { panic("must not be called") }
+func (k noKey) DecodeNonTerminal(_ []byte) (int, noKey, error)   { panic("must not be called") }
+func (k noKey) SizeNonTerminal(_ noKey) int                      { panic("must not be called") }

collections/iter.go

@@ -20,80 +20,96 @@ const (
 	OrderDescending Order = 1
 )
 
-// BoundInclusive creates a Bound of the provided key K
-// which is inclusive. Meaning, if it is used as Ranger.RangeValues start,
-// the provided key will be included if it exists in the Iterator range.
-func BoundInclusive[K any](key K) *Bound[K] {
-	return &Bound[K]{
-		value:     key,
-		inclusive: true,
-	}
+type rangeKeyKind uint8
+
+const (
+	rangeKeyExact rangeKeyKind = iota
+	rangeKeyNext
+	rangeKeyPrefixEnd
+)
+
+// RangeKey wraps a generic range key K, acts as an enum which defines different
+// ways to encode the wrapped key to bytes when it's being used in an iteration.
+type RangeKey[K any] struct {
+	kind rangeKeyKind
+	key  K
 }
 
-// BoundExclusive creates a Bound of the provided key K
-// which is exclusive. Meaning, if it is used as Ranger.RangeValues start,
-// the provided key will be excluded if it exists in the Iterator range.
-func BoundExclusive[K any](key K) *Bound[K] {
-	return &Bound[K]{
-		value:     key,
-		inclusive: false,
-	}
+// RangeKeyNext instantiates a RangeKey that when encoded to bytes
+// identifies the next key after the provided key K.
+// Example: given a string key "ABCD" the next key is bytes("ABCD\0")
+// It's useful when defining inclusivity or exclusivity of a key
+// in store iteration. Specifically: to make an Iterator start exclude key K
+// I would return a RangeKeyNext(key) in the Ranger start.
+func RangeKeyNext[K any](key K) *RangeKey[K] {
+	return &RangeKey[K]{key: key, kind: rangeKeyNext}
 }
 
-// Bound defines key bounds for Start and Ends of iterator ranges.
-type Bound[K any] struct {
-	value     K
-	inclusive bool
+// RangeKeyPrefixEnd instantiates a RangeKey that when encoded to bytes
+// identifies the key that would end the prefix of the key K.
+// Example: if the string key "ABCD" is provided, it would be encoded as bytes("ABCE").
+func RangeKeyPrefixEnd[K any](key K) *RangeKey[K] {
+	return &RangeKey[K]{key: key, kind: rangeKeyPrefixEnd}
+}
+
+// RangeKeyExact instantiates a RangeKey that applies no modifications
+// to the key K. So its bytes representation will not be altered.
+func RangeKeyExact[K any](key K) *RangeKey[K] {
+	return &RangeKey[K]{key: key, kind: rangeKeyExact}
 }
 
 // Ranger defines a generic interface that provides a range of keys.
 type Ranger[K any] interface {
 	// RangeValues is defined by Ranger implementers.
-	// It provides instructions to generate an Iterator instance.
-	// If prefix is not nil, then the Iterator will return only the keys which start
-	// with the given prefix.
-	// If start is not nil, then the Iterator will return only keys which are greater than the provided start
-	// or greater equal depending on the bound is inclusive or exclusive.
-	// If end is not nil, then the Iterator will return only keys which are smaller than the provided end
-	// or smaller equal depending on the bound is inclusive or exclusive.
-	RangeValues() (prefix *K, start *Bound[K], end *Bound[K], order Order, err error)
+	// The implementer can optionally return a start and an end.
+	// If start is nil and end is not, the iteration will include all the keys
+	// in the collection up until the provided end.
+	// If start is defined and end is nil, the iteration will include all the keys
+	// in the collection starting from the provided start.
+	// If both are nil then the iteration will include all the possible keys in the
+	// collection.
+	// Order defines the order of the iteration, if order is OrderAscending then the
+	// iteration will yield keys from the smallest to the biggest, if order
+	// is OrderDescending then the iteration will yield keys from the biggest to the smallest.
+	// Ordering is defined by the keys bytes representation, which is dependent on the KeyCodec used.
+	RangeValues() (start *RangeKey[K], end *RangeKey[K], order Order, err error)
 }
 
 // Range is a Ranger implementer.
 type Range[K any] struct {
-	prefix *K
-	start  *Bound[K]
-	end    *Bound[K]
-	order  Order
+	start *RangeKey[K]
+	end   *RangeKey[K]
+	order Order
 }
 
 // Prefix sets a fixed prefix for the key range.
 func (r *Range[K]) Prefix(key K) *Range[K] {
-	r.prefix = &key
+	r.start = RangeKeyExact(key)
+	r.end = RangeKeyPrefixEnd(key)
 	return r
 }
 
 // StartInclusive makes the range contain only keys which are bigger or equal to the provided start K.
 func (r *Range[K]) StartInclusive(start K) *Range[K] {
-	r.start = BoundInclusive(start)
+	r.start = RangeKeyExact(start)
 	return r
 }
 
 // StartExclusive makes the range contain only keys which are bigger to the provided start K.
 func (r *Range[K]) StartExclusive(start K) *Range[K] {
-	r.start = BoundExclusive(start)
+	r.start = RangeKeyNext(start)
 	return r
 }
 
 // EndInclusive makes the range contain only keys which are smaller or equal to the provided end K.
 func (r *Range[K]) EndInclusive(end K) *Range[K] {
-	r.end = BoundInclusive(end)
+	r.end = RangeKeyNext(end)
 	return r
 }
 
 // EndExclusive makes the range contain only keys which are smaller to the provided end K.
 func (r *Range[K]) EndExclusive(end K) *Range[K] {
-	r.end = BoundExclusive(end)
+	r.end = RangeKeyExact(end)
 	return r
 }
 
@@ -108,118 +124,64 @@ var (
 	errOrder = errors.New("collections: invalid order")
 )
 
-func (r *Range[K]) RangeValues() (prefix *K, start *Bound[K], end *Bound[K], order Order, err error) {
-	if r.prefix != nil && (r.end != nil || r.start != nil) {
-		return nil, nil, nil, order, fmt.Errorf("%w: prefix must not be set if either start or end are specified", errRange)
-	}
-	return r.prefix, r.start, r.end, r.order, nil
+func (r *Range[K]) RangeValues() (start *RangeKey[K], end *RangeKey[K], order Order, err error) {
+	return r.start, r.end, r.order, nil
 }
 
 // iteratorFromRanger generates an Iterator instance, with the proper prefixing and ranging.
 // a nil Ranger can be seen as an ascending iteration over all the possible keys.
 func iteratorFromRanger[K, V any](ctx context.Context, m Map[K, V], r Ranger[K]) (iter Iterator[K, V], err error) {
 	var (
-		prefix *K
-		start  *Bound[K]
-		end    *Bound[K]
-		order  = OrderAscending
+		start *RangeKey[K]
+		end   *RangeKey[K]
+		order = OrderAscending
 	)
-	// if Ranger is specified then we override the defaults
+
 	if r != nil {
-		prefix, start, end, order, err = r.RangeValues()
+		start, end, order, err = r.RangeValues()
 		if err != nil {
 			return iter, err
 		}
 	}
-	if prefix != nil && (start != nil || end != nil) {
-		return iter, fmt.Errorf("%w: prefix must not be set if either start or end are specified", errRange)
-	}
 
-	// compute start and end bytes
-	var startBytes, endBytes []byte
-	if prefix != nil {
-		startBytes, endBytes, err = prefixStartEndBytes(m, *prefix)
+	startBytes := m.prefix
+	if start != nil {
+		startBytes, err = encodeRangeBound(m.prefix, m.kc, start)
 		if err != nil {
 			return iter, err
 		}
-	} else {
-		startBytes, endBytes, err = rangeStartEndBytes(m, start, end)
+	}
+	var endBytes []byte
+	if end != nil {
+		endBytes, err = encodeRangeBound(m.prefix, m.kc, end)
 		if err != nil {
 			return iter, err
 		}
+	} else {
+		endBytes = nextBytesPrefixKey(m.prefix)
 	}
 
-	// get store
 	kv := m.sa(ctx)
-
-	// create iter
-	var storeIter store.Iterator
 	switch order {
 	case OrderAscending:
-		storeIter = kv.Iterator(startBytes, endBytes)
+		return newIterator(kv.Iterator(startBytes, endBytes), m)
 	case OrderDescending:
-		storeIter = kv.ReverseIterator(startBytes, endBytes)
+		return newIterator(kv.ReverseIterator(startBytes, endBytes), m)
 	default:
-		return iter, fmt.Errorf("%w: %d", errOrder, order)
-	}
-
-	// check if valid
-	if !storeIter.Valid() {
-		return iter, ErrInvalidIterator
+		return iter, errOrder
 	}
-
-	// all good
-	iter.kc = m.kc
-	iter.vc = m.vc
-	iter.prefixLength = len(m.prefix)
-	iter.iter = storeIter
-	return iter, nil
 }
 
-// rangeStartEndBytes computes a range's start and end bytes to be passed to the store's iterator.
-func rangeStartEndBytes[K, V any](m Map[K, V], start, end *Bound[K]) (startBytes, endBytes []byte, err error) {
-	startBytes = m.prefix
-	if start != nil {
-		startBytes, err = encodeKeyWithPrefix(m.prefix, m.kc, start.value)
-		if err != nil {
-			return startBytes, endBytes, err
-		}
-		// the start of iterators is by default inclusive,
-		// in order to make it exclusive we extend the start
-		// by one single byte.
-		if !start.inclusive {
-			startBytes = extendOneByte(startBytes)
-		}
-	}
-	if end != nil {
-		endBytes, err = encodeKeyWithPrefix(m.prefix, m.kc, end.value)
-		if err != nil {
-			return startBytes, endBytes, err
-		}
-		// the end of iterators is by default exclusive
-		// in order to make it inclusive we extend the end
-		// by one single byte.
-		if end.inclusive {
-			endBytes = extendOneByte(endBytes)
-		}
-	} else {
-		// if end is not specified then we simply are
-		// inclusive up to the last key of the Prefix
-		// of the collection.
-		endBytes = prefixEndBytes(m.prefix)
-	}
-
-	return startBytes, endBytes, nil
-}
-
-// prefixStartEndBytes returns the start and end bytes to be provided to the store's iterator, considering we're prefixing
-// over a specific key.
-func prefixStartEndBytes[K, V any](m Map[K, V], prefix K) (startBytes, endBytes []byte, err error) {
-	startBytes, err = encodeKeyWithPrefix(m.prefix, m.kc, prefix)
-	if err != nil {
-		return
+func newIterator[K, V any](iterator store.Iterator, m Map[K, V]) (Iterator[K, V], error) {
+	if iterator.Valid() == false {
+		return Iterator[K, V]{}, ErrInvalidIterator
 	}
-	return startBytes, prefixEndBytes(startBytes), nil
+	return Iterator[K, V]{
+		kc:           m.kc,
+		vc:           m.vc,
+		iter:         iterator,
+		prefixLength: len(m.prefix),
+	}, nil
 }
 
 // Iterator defines a generic wrapper around an storetypes.Iterator.
@@ -327,14 +289,33 @@ type KeyValue[K, V any] struct {
 	Value V
 }
 
-func extendOneByte(b []byte) []byte {
+// encodeRangeBound encodes a range bound, modifying the key bytes to adhere to bound semantics.
+func encodeRangeBound[T any](prefix []byte, keyCodec KeyCodec[T], bound *RangeKey[T]) ([]byte, error) {
+	key, err := encodeKeyWithPrefix(prefix, keyCodec, bound.key)
+	if err != nil {
+		return nil, err
+	}
+	switch bound.kind {
+	case rangeKeyExact:
+		return key, nil
+	case rangeKeyNext:
+		return nextBytesKey(key), nil
+	case rangeKeyPrefixEnd:
+		return nextBytesPrefixKey(key), nil
+	default:
+		panic("undefined bound kind")
+	}
+}
+
+// nextBytesKey returns the next byte key after this one.
+func nextBytesKey(b []byte) []byte {
 	return append(b, 0)
 }
 
-// prefixEndBytes returns the []byte that would end a
+// nextBytesPrefixKey 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 {
+func nextBytesPrefixKey(prefix []byte) []byte {
 	if len(prefix) == 0 {
 		return nil
 	}