feat(collections): add alternative value codec (#16773)

Co-authored-by: unknown unknown <unknown@unknown>

collections/CHANGELOG.md

@@ -34,6 +34,9 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### Features
 
 * [#16074](https://github.com/cosmos/cosmos-sdk/pull/16607) - Introduces `Clear` method for `Map` and `KeySet`
+* [#16773](https://github.com/cosmos/cosmos-sdk/pull/16773)
+  * Adds `AltValueCodec` which provides a way to decode a value in two ways.
+  * Adds the possibility to specify an alternative way to decode the values of `KeySet`, `indexes.Multi`, `indexes.ReversePair`. 
 
 ## [v0.2.0](https://github.com/cosmos/cosmos-sdk/releases/tag/collections%2Fv0.2.0)
 

collections/README.md

@@ -1117,3 +1117,30 @@ func (k Keeper) GetAccount(ctx sdk.context, addr sdk.AccAddress) (sdk.AccountI,
     return k.Accounts.Get(ctx, addr)
 }
 ```
+
+## Advanced Usages
+
+### Alternative Value Codec
+
+The `codec.AltValueCodec` allows a collection to decode values using a different codec than the one used to encode them.
+Basically it enables to decode two different byte representations of the same concrete value.
+It can be used to lazily migrate values from one bytes representation to another, as long as the new representation is
+not able to decode the old one.
+
+A concrete example can be found in `x/bank` where the balance was initially stored as `Coin` and then migrated to `Int`.
+
+```go
+
+var BankBalanceValueCodec = codec.NewAltValueCodec(sdk.IntValue, func(b []byte) (sdk.Int, error) {
+    coin := sdk.Coin{}
+    err := coin.Unmarshal(b)
+    if err != nil {
+        return sdk.Int{}, err
+    }
+    return coin.Amount, nil
+})
+```
+
+The above example shows how to create an `AltValueCodec` that can decode both `sdk.Int` and `sdk.Coin` values. The provided 
+decoder function will be used as a fallback in case the default decoder fails. When the value will be encoded back into state
+it will use the default encoder. This allows to lazily migrate values to a new bytes representation.

collections/codec/alternative_value.go

@@ -0,0 +1,47 @@
+package codec
+
+// NewAltValueCodec returns a new AltValueCodec. canonicalValueCodec is the codec that you want the value
+// to be encoded and decoded as, alternativeDecoder is a function that will attempt to decode the value
+// in case the canonicalValueCodec fails to decode it.
+func NewAltValueCodec[V any](canonicalValueCodec ValueCodec[V], alternativeDecoder func([]byte) (V, error)) ValueCodec[V] {
+	return AltValueCodec[V]{
+		canonicalValueCodec: canonicalValueCodec,
+		alternativeDecoder:  alternativeDecoder,
+	}
+}
+
+// AltValueCodec is a codec that can decode a value from state in an alternative format.
+// This is useful for migrating data from one format to another. For example, in x/bank
+// balances were initially encoded as sdk.Coin, now they are encoded as math.Int.
+// The AltValueCodec will be trying to decode the value as math.Int, and if that fails,
+// it will attempt to decode it as sdk.Coin.
+// NOTE: if the canonical format can also decode the alternative format, then this codec
+// will produce undefined and undesirable behavior.
+type AltValueCodec[V any] struct {
+	canonicalValueCodec ValueCodec[V]
+	alternativeDecoder  func([]byte) (V, error)
+}
+
+// Decode will attempt to decode the value from state using the canonical value codec.
+// If it fails to decode, it will attempt to decode the value using the alternative decoder.
+func (a AltValueCodec[V]) Decode(b []byte) (V, error) {
+	v, err := a.canonicalValueCodec.Decode(b)
+	if err != nil {
+		return a.alternativeDecoder(b)
+	}
+	return v, nil
+}
+
+// Below there is the implementation of ValueCodec relying on the canonical value codec.
+
+func (a AltValueCodec[V]) Encode(value V) ([]byte, error) { return a.canonicalValueCodec.Encode(value) }
+
+func (a AltValueCodec[V]) EncodeJSON(value V) ([]byte, error) {
+	return a.canonicalValueCodec.EncodeJSON(value)
+}
+
+func (a AltValueCodec[V]) DecodeJSON(b []byte) (V, error) { return a.canonicalValueCodec.DecodeJSON(b) }
+
+func (a AltValueCodec[V]) Stringify(value V) string { return a.canonicalValueCodec.Stringify(value) }
+
+func (a AltValueCodec[V]) ValueType() string { return a.canonicalValueCodec.ValueType() }

collections/codec/alternative_value_test.go

@@ -0,0 +1,53 @@
+package codec_test
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"cosmossdk.io/collections/codec"
+	"cosmossdk.io/collections/colltest"
+)
+
+type altValue struct {
+	Value uint64 `json:"value"`
+}
+
+func TestAltValueCodec(t *testing.T) {
+	// we assume we want to migrate the value from json(altValue) to just be
+	// the raw value uint64.
+	canonical := codec.KeyToValueCodec(codec.NewUint64Key[uint64]())
+	alternative := func(v []byte) (uint64, error) {
+		var alt altValue
+		err := json.Unmarshal(v, &alt)
+		if err != nil {
+			return 0, err
+		}
+		return alt.Value, nil
+	}
+
+	cdc := codec.NewAltValueCodec(canonical, alternative)
+
+	t.Run("decodes alternative value", func(t *testing.T) {
+		expected := uint64(100)
+		alternativeEncodedBytes, err := json.Marshal(altValue{Value: expected})
+		require.NoError(t, err)
+		got, err := cdc.Decode(alternativeEncodedBytes)
+		require.NoError(t, err)
+		require.Equal(t, expected, got)
+	})
+
+	t.Run("decodes canonical value", func(t *testing.T) {
+		expected := uint64(100)
+		canonicalEncodedBytes, err := cdc.Encode(expected)
+		require.NoError(t, err)
+		got, err := cdc.Decode(canonicalEncodedBytes)
+		require.NoError(t, err)
+		require.Equal(t, expected, got)
+	})
+
+	t.Run("conformance", func(t *testing.T) {
+		colltest.TestValueCodec(t, cdc, uint64(100))
+	})
+}

collections/codec/bool.go

@@ -33,7 +33,7 @@ func (b boolKey[T]) Decode(buffer []byte) (int, T, error) {
 	}
 }
 
-func (b boolKey[T]) Size(key T) int { return 1 }
+func (b boolKey[T]) Size(_ T) int { return 1 }
 
 func (b boolKey[T]) EncodeJSON(value T) ([]byte, error) {
 	return json.Marshal(value)

collections/colltest/codec.go

@@ -40,6 +40,11 @@ func TestKeyCodec[T any](t *testing.T, keyCodec codec.KeyCodec[T], key T) {
 	decoded, err := keyCodec.DecodeJSON(keyJSON)
 	require.NoError(t, err)
 	require.Equal(t, key, decoded, "json encoding and decoding did not produce the same results")
+
+	// check type
+	require.NotEmpty(t, keyCodec.KeyType())
+	// check string
+	_ = keyCodec.Stringify(key)
 }
 
 // TestValueCodec asserts the correct behavior of a ValueCodec over the type T.

collections/indexes/multi.go

@@ -8,6 +8,21 @@ import (
 	"cosmossdk.io/collections/codec"
 )
 
+type multiOptions struct {
+	uncheckedValue bool
+}
+
+// WithMultiUncheckedValue is an option that can be passed to NewMulti to
+// ignore index values different from '[]byte{}' and continue with the operation.
+// This should be used only to behave nicely in case you have used values different
+// from '[]byte{}' in your storage before migrating to collections. Refer to
+// WithKeySetUncheckedValue for more information.
+func WithMultiUncheckedValue() func(*multiOptions) {
+	return func(o *multiOptions) {
+		o.uncheckedValue = true
+	}
+}
+
 // Multi defines the most common index. It can be used to create a reference between
 // a field of value and its primary key. Multiple primary keys can be mapped to the same
 // reference key as the index does not enforce uniqueness constraints.
@@ -27,7 +42,19 @@ func NewMulti[ReferenceKey, PrimaryKey, Value any](
 	refCodec codec.KeyCodec[ReferenceKey],
 	pkCodec codec.KeyCodec[PrimaryKey],
 	getRefKeyFunc func(pk PrimaryKey, value Value) (ReferenceKey, error),
+	options ...func(*multiOptions),
 ) *Multi[ReferenceKey, PrimaryKey, Value] {
+	o := new(multiOptions)
+	for _, opt := range options {
+		opt(o)
+	}
+	if o.uncheckedValue {
+		return &Multi[ReferenceKey, PrimaryKey, Value]{
+			getRefKey: getRefKeyFunc,
+			refKeys:   collections.NewKeySet(schema, prefix, name, collections.PairKeyCodec(refCodec, pkCodec), collections.WithKeySetUncheckedValue()),
+		}
+	}
+
 	return &Multi[ReferenceKey, PrimaryKey, Value]{
 		getRefKey: getRefKeyFunc,
 		refKeys:   collections.NewKeySet(schema, prefix, name, collections.PairKeyCodec(refCodec, pkCodec)),

collections/indexes/multi_test.go

@@ -61,3 +61,52 @@ func TestMultiIndex(t *testing.T) {
 	require.False(t, iter.Valid())
 	require.NoError(t, iter.Close())
 }
+
+func TestMultiUnchecked(t *testing.T) {
+	sk, ctx := deps()
+	schema := collections.NewSchemaBuilder(sk)
+
+	uncheckedMi := NewMulti(schema, collections.NewPrefix("prefix"), "multi_index", collections.StringKey, collections.Uint64Key, func(_ uint64, value company) (string, error) {
+		return value.City, nil
+	}, WithMultiUncheckedValue())
+
+	mi := NewMulti(schema, collections.NewPrefix("prefix"), "multi_index", collections.StringKey, collections.Uint64Key, func(_ uint64, value company) (string, error) {
+		return value.City, nil
+	})
+
+	rawKey, err := collections.EncodeKeyWithPrefix(
+		collections.NewPrefix("prefix"),
+		uncheckedMi.KeyCodec(),
+		collections.Join("milan", uint64(2)))
+	require.NoError(t, err)
+
+	// set value to be something different from []byte{}
+	require.NoError(t, sk.OpenKVStore(ctx).Set(rawKey, []byte("something")))
+
+	// normal multi index will fail.
+	err = mi.Walk(ctx, nil, func(indexingKey string, indexedKey uint64) (stop bool, err error) {
+		return true, err
+	})
+	require.ErrorIs(t, err, collections.ErrEncoding)
+
+	// unchecked multi index will not fail.
+	err = uncheckedMi.Walk(ctx, nil, func(indexingKey string, indexedKey uint64) (stop bool, err error) {
+		require.Equal(t, "milan", indexingKey)
+		require.Equal(t, uint64(2), indexedKey)
+		return true, err
+	})
+	require.NoError(t, err)
+
+	// unchecked multi will also reset the value
+	err = mi.Reference(ctx, 2, company{City: "milan"}, func() (company, error) {
+		return company{
+			City: "milan",
+		}, nil
+	})
+	require.NoError(t, err)
+
+	// value reset to []byte{}
+	rawValue, err := sk.OpenKVStore(ctx).Get(rawKey)
+	require.NoError(t, err)
+	require.Equal(t, []byte{}, rawValue)
+}

collections/indexes/reverse_pair.go

@@ -7,6 +7,21 @@ import (
 	"cosmossdk.io/collections/codec"
 )
 
+type reversePairOptions struct {
+	uncheckedValue bool
+}
+
+// WithReversePairUncheckedValue is an option that can be passed to NewReversePair to
+// ignore index values different from '[]byte{}' and continue with the operation.
+// This should be used only if you are migrating to collections and have used a different
+// placeholder value in your storage index keys.
+// Refer to WithKeySetUncheckedValue for more information.
+func WithReversePairUncheckedValue() func(*reversePairOptions) {
+	return func(o *reversePairOptions) {
+		o.uncheckedValue = true
+	}
+}
+
 // ReversePair is an index that is used with collections.Pair keys. It indexes objects by their second part of the key.
 // When the value is being indexed by collections.IndexedMap then ReversePair will create a relationship between
 // the second part of the primary key and the first part.
@@ -31,8 +46,19 @@ func NewReversePair[Value, K1, K2 any](
 	prefix collections.Prefix,
 	name string,
 	pairCodec codec.KeyCodec[collections.Pair[K1, K2]],
+	options ...func(*reversePairOptions),
 ) *ReversePair[K1, K2, Value] {
 	pkc := pairCodec.(pairKeyCodec[K1, K2])
+	o := new(reversePairOptions)
+	for _, option := range options {
+		option(o)
+	}
+	if o.uncheckedValue {
+		return &ReversePair[K1, K2, Value]{
+			refKeys: collections.NewKeySet(sb, prefix, name, collections.PairKeyCodec(pkc.KeyCodec2(), pkc.KeyCodec1()), collections.WithKeySetUncheckedValue()),
+		}
+	}
+
 	mi := &ReversePair[K1, K2, Value]{
 		refKeys: collections.NewKeySet(sb, prefix, name, collections.PairKeyCodec(pkc.KeyCodec2(), pkc.KeyCodec1())),
 	}

collections/indexes/reverse_pair_test.go

@@ -67,3 +67,38 @@ func TestReversePair(t *testing.T) {
 	_, err = indexedMap.Indexes.Denom.MatchExact(ctx, "atom")
 	require.ErrorIs(t, collections.ErrInvalidIterator, err)
 }
+
+func TestUncheckedReversePair(t *testing.T) {
+	sk, ctx := deps()
+	sb := collections.NewSchemaBuilder(sk)
+	prefix := collections.NewPrefix("prefix")
+	keyCodec := collections.PairKeyCodec(collections.StringKey, collections.StringKey)
+
+	uncheckedRp := NewReversePair[Amount](sb, prefix, "denom_index", keyCodec, WithReversePairUncheckedValue())
+	rp := NewReversePair[Amount](sb, prefix, "denom_index", keyCodec)
+
+	rawKey, err := collections.EncodeKeyWithPrefix(prefix, uncheckedRp.KeyCodec(), collections.Join("atom", "address1"))
+	require.NoError(t, err)
+
+	require.NoError(t, sk.OpenKVStore(ctx).Set(rawKey, []byte("i should not be here")))
+
+	// normal reverse pair fails
+	err = rp.Walk(ctx, nil, func(denom, address string) (bool, error) {
+		return false, nil
+	})
+	require.ErrorIs(t, err, collections.ErrEncoding)
+
+	// unchecked reverse pair succeeds
+	err = uncheckedRp.Walk(ctx, nil, func(indexingKey, indexedKey string) (stop bool, err error) {
+		require.Equal(t, "atom", indexingKey)
+		return true, nil
+	})
+	require.NoError(t, err)
+
+	// unchecked reverse pair lazily updates
+	err = uncheckedRp.Reference(ctx, collections.Join("address1", "atom"), 0, nil)
+	require.NoError(t, err)
+	rawValue, err := sk.OpenKVStore(ctx).Get(rawKey)
+	require.NoError(t, err)
+	require.Equal(t, []byte{}, rawValue)
+}

collections/keyset.go

@@ -8,14 +8,43 @@ import (
 	"cosmossdk.io/collections/codec"
 )
 
+// WithKeySetUncheckedValue changes the behavior of the KeySet when it encounters
+// a value different from '[]byte{}', by default the KeySet errors when this happens.
+// This option allows to ignore the value and continue with the operation, in turn
+// the value will be cleared out and set to '[]byte{}'.
+// You should never use this option if you're creating a new state object from scratch.
+// This should be used only to behave nicely in case you have used values different
+// from '[]byte{}' in your storage before migrating to collections.
+func WithKeySetUncheckedValue() func(opt *keySetOptions) {
+	return func(opt *keySetOptions) {
+		opt.uncheckedValue = true
+	}
+}
+
+type keySetOptions struct{ uncheckedValue bool }
+
 // KeySet builds on top of a Map and represents a collection retaining only a set
 // of keys and no value. It can be used, for example, in an allow list.
 type KeySet[K any] Map[K, NoValue]
 
 // NewKeySet returns a KeySet given a Schema, Prefix a human name for the collection
 // and a KeyCodec for the key K.
-func NewKeySet[K any](schema *SchemaBuilder, prefix Prefix, name string, keyCodec codec.KeyCodec[K]) KeySet[K] {
-	return (KeySet[K])(NewMap(schema, prefix, name, keyCodec, noValueCodec))
+func NewKeySet[K any](
+	schema *SchemaBuilder,
+	prefix Prefix,
+	name string,
+	keyCodec codec.KeyCodec[K],
+	options ...func(opt *keySetOptions),
+) KeySet[K] {
+	o := new(keySetOptions)
+	for _, opt := range options {
+		opt(o)
+	}
+	vc := noValueCodec
+	if o.uncheckedValue {
+		vc = codec.NewAltValueCodec(vc, func(_ []byte) (NoValue, error) { return NoValue{}, nil })
+	}
+	return (KeySet[K])(NewMap(schema, prefix, name, keyCodec, vc))
 }
 
 // Set adds the key to the KeySet. Errors on encoding problems.
@@ -79,6 +108,7 @@ var noValueCodec codec.ValueCodec[NoValue] = NoValue{}
 
 const noValueValueType = "no_value"
 
+// NoValue is a type that can be used to represent a non-existing value.
 type NoValue struct{}
 
 func (n NoValue) EncodeJSON(_ NoValue) ([]byte, error) {
@@ -98,7 +128,7 @@ func (NoValue) Encode(_ NoValue) ([]byte, error) {
 
 func (NoValue) Decode(b []byte) (NoValue, error) {
 	if !bytes.Equal(b, []byte{}) {
-		return NoValue{}, fmt.Errorf("%w: invalid value, wanted an empty non-nil byte slice", ErrEncoding)
+		return NoValue{}, fmt.Errorf("%w: invalid value, wanted an empty non-nil byte slice, got: %x", ErrEncoding, b)
 	}
 	return NoValue{}, nil
 }

collections/keyset_test.go

@@ -67,3 +67,34 @@ func Test_noValue(t *testing.T) {
 	_, err = noValueCodec.Decode([]byte("bad"))
 	require.ErrorIs(t, err, ErrEncoding)
 }
+
+func TestUncheckedKeySet(t *testing.T) {
+	sk, ctx := deps()
+	schema := NewSchemaBuilder(sk)
+	uncheckedKs := NewKeySet(schema, NewPrefix("keyset"), "keyset", StringKey, WithKeySetUncheckedValue())
+	ks := NewKeySet(schema, NewPrefix("keyset"), "keyset", StringKey)
+	// we set a NoValue unfriendly value.
+	require.NoError(t, sk.OpenKVStore(ctx).Set([]byte("keyset1"), []byte("A")))
+	require.NoError(t, sk.OpenKVStore(ctx).Set([]byte("keyset2"), []byte("B")))
+
+	// the standard KeySet errors here, because it doesn't like the fact that the value is []byte("A")
+	// and not []byte{}.
+	err := ks.Walk(ctx, nil, func(key string) (stop bool, err error) {
+		return true, nil
+	})
+	require.ErrorIs(t, err, ErrEncoding)
+
+	// the unchecked KeySet doesn't care about the value, so it works.
+	err = uncheckedKs.Walk(ctx, nil, func(key string) (stop bool, err error) {
+		require.Equal(t, "1", key)
+		return true, nil
+	})
+	require.NoError(t, err)
+
+	// now we set it again
+	require.NoError(t, uncheckedKs.Set(ctx, "1"))
+	// and we will see that the value which was []byte("A") has been cleared to be []byte{}
+	raw, err := sk.OpenKVStore(ctx).Get([]byte("keyset1"))
+	require.NoError(t, err)
+	require.Equal(t, []byte{}, raw)
+}