feat(collections): add optional key and value naming methods (#20538)

collections/CHANGELOG.md

@@ -39,6 +39,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#19861](https://github.com/cosmos/cosmos-sdk/pull/19861) Add `NewJSONValueCodec` value codec as an alternative for `codec.CollValue` from the SDK for non protobuf types.
 * [#21090](https://github.com/cosmos/cosmos-sdk/pull/21090) Introduces `Quad`, a composite key with four keys.
 * [#20704](https://github.com/cosmos/cosmos-sdk/pull/20704) Add `ModuleCodec` method to `Schema` and `HasSchemaCodec` interface in order to support `cosmossdk.io/schema` compatible indexing.
+* [#20538](https://github.com/cosmos/cosmos-sdk/pull/20538) Add `Nameable` variations to `KeyCodec` and `ValueCodec` to allow for better indexing of `collections` types.
 
 ## [v0.4.0](https://github.com/cosmos/cosmos-sdk/releases/tag/collections%2Fv0.4.0)
 

collections/codec/bool.go

@@ -6,7 +6,7 @@ import (
 	"strconv"
 )
 
-func NewBoolKey[T ~bool]() KeyCodec[T] { return boolKey[T]{} }
+func NewBoolKey[T ~bool]() NameableKeyCodec[T] { return boolKey[T]{} }
 
 type boolKey[T ~bool] struct{}
 
@@ -64,3 +64,7 @@ func (b boolKey[T]) DecodeNonTerminal(buffer []byte) (int, T, error) {
 func (b boolKey[T]) SizeNonTerminal(key T) int {
 	return b.Size(key)
 }
+
+func (b boolKey[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: b, Name: name}
+}

collections/codec/bytes.go

@@ -10,7 +10,7 @@ import (
 // using the BytesKey KeyCodec.
 const MaxBytesKeyNonTerminalSize = math.MaxUint8
 
-func NewBytesKey[T ~[]byte]() KeyCodec[T] { return bytesKey[T]{} }
+func NewBytesKey[T ~[]byte]() NameableKeyCodec[T] { return bytesKey[T]{} }
 
 type bytesKey[T ~[]byte] struct{}
 
@@ -77,3 +77,7 @@ func (bytesKey[T]) DecodeNonTerminal(buffer []byte) (int, T, error) {
 func (bytesKey[T]) SizeNonTerminal(key T) int {
 	return len(key) + 1
 }
+
+func (b bytesKey[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: b, Name: name}
+}

collections/codec/codec.go

@@ -125,7 +125,9 @@ type UntypedValueCodec struct {
 }
 
 // KeyToValueCodec converts a KeyCodec into a ValueCodec.
-func KeyToValueCodec[K any](keyCodec KeyCodec[K]) ValueCodec[K] { return keyToValueCodec[K]{keyCodec} }
+func KeyToValueCodec[K any](keyCodec KeyCodec[K]) NameableValueCodec[K] {
+	return keyToValueCodec[K]{kc: keyCodec}
+}
 
 // keyToValueCodec is a ValueCodec that wraps a KeyCodec to make it behave like a ValueCodec.
 type keyToValueCodec[K any] struct {
@@ -167,3 +169,7 @@ func (k keyToValueCodec[K]) Stringify(value K) string {
 func (k keyToValueCodec[K]) ValueType() string {
 	return k.kc.KeyType()
 }
+
+func (k keyToValueCodec[K]) WithName(name string) ValueCodec[K] {
+	return NamedValueCodec[K]{ValueCodec: k, Name: name}
+}

collections/codec/codec_test.go

@@ -7,7 +7,7 @@ import (
 )
 
 func TestUntypedValueCodec(t *testing.T) {
-	vc := NewUntypedValueCodec(KeyToValueCodec(NewStringKeyCodec[string]()))
+	vc := NewUntypedValueCodec(ValueCodec[string](KeyToValueCodec(KeyCodec[string](NewStringKeyCodec[string]()))))
 
 	t.Run("encode/decode", func(t *testing.T) {
 		_, err := vc.Encode(0)

collections/codec/int.go

@@ -7,7 +7,7 @@ import (
 	"strconv"
 )
 
-func NewInt64Key[T ~int64]() KeyCodec[T] { return int64Key[T]{} }
+func NewInt64Key[T ~int64]() NameableKeyCodec[T] { return int64Key[T]{} }
 
 type int64Key[T ~int64] struct{}
 
@@ -64,7 +64,11 @@ func (i int64Key[T]) SizeNonTerminal(_ T) int {
 	return 8
 }
 
-func NewInt32Key[T ~int32]() KeyCodec[T] {
+func (i int64Key[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: i, Name: name}
+}
+
+func NewInt32Key[T ~int32]() NameableKeyCodec[T] {
 	return int32Key[T]{}
 }
 
@@ -121,3 +125,7 @@ func (i int32Key[T]) DecodeNonTerminal(buffer []byte) (int, T, error) {
 func (i int32Key[T]) SizeNonTerminal(_ T) int {
 	return 4
 }
+
+func (i int32Key[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: i, Name: name}
+}

collections/codec/naming.go

@@ -0,0 +1,63 @@
+package codec
+
+import "fmt"
+
+// NameableKeyCodec is a KeyCodec that can be named.
+type NameableKeyCodec[T any] interface {
+	KeyCodec[T]
+
+	// WithName returns the KeyCodec with the provided name.
+	WithName(name string) KeyCodec[T]
+}
+
+// NameableValueCodec is a ValueCodec that can be named.
+type NameableValueCodec[T any] interface {
+	ValueCodec[T]
+
+	// WithName returns the ValueCodec with the provided name.
+	WithName(name string) ValueCodec[T]
+}
+
+// NamedKeyCodec wraps a KeyCodec with a name.
+// The underlying key codec MUST have exactly one field in its schema.
+type NamedKeyCodec[T any] struct {
+	KeyCodec[T]
+
+	// Name is the name of the KeyCodec in the schema.
+	Name string
+}
+
+// SchemaCodec returns the schema codec for the named key codec.
+func (n NamedKeyCodec[T]) SchemaCodec() (SchemaCodec[T], error) {
+	cdc, err := KeySchemaCodec[T](n.KeyCodec)
+	if err != nil {
+		return SchemaCodec[T]{}, err
+	}
+	return withName(cdc, n.Name)
+}
+
+// NamedValueCodec wraps a ValueCodec with a name.
+// The underlying value codec MUST have exactly one field in its schema.
+type NamedValueCodec[T any] struct {
+	ValueCodec[T]
+
+	// Name is the name of the ValueCodec in the schema.
+	Name string
+}
+
+// SchemaCodec returns the schema codec for the named value codec.
+func (n NamedValueCodec[T]) SchemaCodec() (SchemaCodec[T], error) {
+	cdc, err := ValueSchemaCodec[T](n.ValueCodec)
+	if err != nil {
+		return SchemaCodec[T]{}, err
+	}
+	return withName(cdc, n.Name)
+}
+
+func withName[T any](cdc SchemaCodec[T], name string) (SchemaCodec[T], error) {
+	if len(cdc.Fields) != 1 {
+		return SchemaCodec[T]{}, fmt.Errorf("expected exactly one field to be named, got %d", len(cdc.Fields))
+	}
+	cdc.Fields[0].Name = name
+	return cdc, nil
+}

collections/codec/string.go

@@ -6,7 +6,7 @@ import (
 	"fmt"
 )
 
-func NewStringKeyCodec[T ~string]() KeyCodec[T] { return stringKey[T]{} }
+func NewStringKeyCodec[T ~string]() NameableKeyCodec[T] { return stringKey[T]{} }
 
 const (
 	// StringDelimiter defines the delimiter of a string key when used in non-terminal encodings.
@@ -66,3 +66,7 @@ func (stringKey[T]) Stringify(key T) string {
 func (stringKey[T]) KeyType() string {
 	return "string"
 }
+
+func (s stringKey[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: s, Name: name}
+}

collections/codec/uint.go

@@ -7,7 +7,7 @@ import (
 	"strconv"
 )
 
-func NewUint64Key[T ~uint64]() KeyCodec[T] { return uint64Key[T]{} }
+func NewUint64Key[T ~uint64]() NameableKeyCodec[T] { return uint64Key[T]{} }
 
 type uint64Key[T ~uint64] struct{}
 
@@ -55,7 +55,11 @@ func (uint64Key[T]) KeyType() string {
 	return "uint64"
 }
 
-func NewUint32Key[T ~uint32]() KeyCodec[T] { return uint32Key[T]{} }
+func (u uint64Key[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: u, Name: name}
+}
+
+func NewUint32Key[T ~uint32]() NameableKeyCodec[T] { return uint32Key[T]{} }
 
 type uint32Key[T ~uint32] struct{}
 
@@ -95,7 +99,11 @@ func (u uint32Key[T]) DecodeNonTerminal(buffer []byte) (int, T, error) { return
 
 func (uint32Key[T]) SizeNonTerminal(_ T) int { return 4 }
 
-func NewUint16Key[T ~uint16]() KeyCodec[T] { return uint16Key[T]{} }
+func (u uint32Key[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: u, Name: name}
+}
+
+func NewUint16Key[T ~uint16]() NameableKeyCodec[T] { return uint16Key[T]{} }
 
 type uint16Key[T ~uint16] struct{}
 
@@ -135,6 +143,10 @@ func (u uint16Key[T]) DecodeNonTerminal(buffer []byte) (int, T, error) { return
 
 func (u uint16Key[T]) SizeNonTerminal(key T) int { return u.Size(key) }
 
+func (u uint16Key[T]) WithName(name string) KeyCodec[T] {
+	return NamedKeyCodec[T]{KeyCodec: u, Name: name}
+}
+
 func uintEncodeJSON(value uint64) ([]byte, error) {
 	str := `"` + strconv.FormatUint(value, 10) + `"`
 	return []byte(str), nil

collections/collections.go

@@ -6,8 +6,9 @@ import (
 	"io"
 	"math"
 
-	"cosmossdk.io/collections/codec"
 	"cosmossdk.io/schema"
+
+	"cosmossdk.io/collections/codec"
 )
 
 var (

collections/naming_test.go

@@ -0,0 +1,56 @@
+package collections
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"cosmossdk.io/collections/codec"
+)
+
+func TestNaming(t *testing.T) {
+	expectKeyCodecName(t, "u16", Uint16Key.WithName("u16"))
+	expectKeyCodecName(t, "u32", Uint32Key.WithName("u32"))
+	expectKeyCodecName(t, "u64", Uint64Key.WithName("u64"))
+	expectKeyCodecName(t, "i32", Int32Key.WithName("i32"))
+	expectKeyCodecName(t, "i64", Int64Key.WithName("i64"))
+	expectKeyCodecName(t, "str", StringKey.WithName("str"))
+	expectKeyCodecName(t, "bytes", BytesKey.WithName("bytes"))
+	expectKeyCodecName(t, "bool", BoolKey.WithName("bool"))
+
+	expectValueCodecName(t, "vu16", Uint16Value.WithName("vu16"))
+	expectValueCodecName(t, "vu32", Uint32Value.WithName("vu32"))
+	expectValueCodecName(t, "vu64", Uint64Value.WithName("vu64"))
+	expectValueCodecName(t, "vi32", Int32Value.WithName("vi32"))
+	expectValueCodecName(t, "vi64", Int64Value.WithName("vi64"))
+	expectValueCodecName(t, "vstr", StringValue.WithName("vstr"))
+	expectValueCodecName(t, "vbytes", BytesValue.WithName("vbytes"))
+	expectValueCodecName(t, "vbool", BoolValue.WithName("vbool"))
+
+	expectKeyCodecNames(t, NamedPairKeyCodec[bool, string]("abc", BoolKey, "def", StringKey), "abc", "def")
+	expectKeyCodecNames(t, NamedTripleKeyCodec[bool, string, int32]("abc", BoolKey, "def", StringKey, "ghi", Int32Key), "abc", "def", "ghi")
+	expectKeyCodecNames(t, NamedQuadKeyCodec[bool, string, int32, uint64]("abc", BoolKey, "def", StringKey, "ghi", Int32Key, "jkl", Uint64Key), "abc", "def", "ghi", "jkl")
+}
+
+func expectKeyCodecName[T any](t *testing.T, name string, cdc codec.KeyCodec[T]) {
+	schema, err := codec.KeySchemaCodec(cdc)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(schema.Fields))
+	require.Equal(t, name, schema.Fields[0].Name)
+}
+
+func expectValueCodecName[T any](t *testing.T, name string, cdc codec.ValueCodec[T]) {
+	schema, err := codec.ValueSchemaCodec(cdc)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(schema.Fields))
+	require.Equal(t, name, schema.Fields[0].Name)
+}
+
+func expectKeyCodecNames[T any](t *testing.T, cdc codec.KeyCodec[T], names ...string) {
+	schema, err := codec.KeySchemaCodec(cdc)
+	require.NoError(t, err)
+	require.Equal(t, len(names), len(schema.Fields))
+	for i, name := range names {
+		require.Equal(t, name, schema.Fields[i].Name)
+	}
+}

collections/pair.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"strings"
 
+	"cosmossdk.io/schema"
+
 	"cosmossdk.io/collections/codec"
 )
 
@@ -54,9 +56,22 @@ func PairKeyCodec[K1, K2 any](keyCodec1 codec.KeyCodec[K1], keyCodec2 codec.KeyC
 	}
 }
 
+// NamedPairKeyCodec instantiates a new KeyCodec instance that can encode the Pair, given the KeyCodec of the
+// first part of the key and the KeyCodec of the second part of the key.
+// It also provides names for the keys which are used for indexing purposes.
+func NamedPairKeyCodec[K1, K2 any](key1Name string, keyCodec1 codec.KeyCodec[K1], key2Name string, keyCodec2 codec.KeyCodec[K2]) codec.KeyCodec[Pair[K1, K2]] {
+	return pairKeyCodec[K1, K2]{
+		key1Name:  key1Name,
+		key2Name:  key2Name,
+		keyCodec1: keyCodec1,
+		keyCodec2: keyCodec2,
+	}
+}
+
 type pairKeyCodec[K1, K2 any] struct {
-	keyCodec1 codec.KeyCodec[K1]
-	keyCodec2 codec.KeyCodec[K2]
+	key1Name, key2Name string
+	keyCodec1          codec.KeyCodec[K1]
+	keyCodec2          codec.KeyCodec[K2]
 }
 
 func (p pairKeyCodec[K1, K2]) KeyCodec1() codec.KeyCodec[K1] { return p.keyCodec1 }
@@ -216,6 +231,39 @@ func (p pairKeyCodec[K1, K2]) DecodeJSON(b []byte) (Pair[K1, K2], error) {
 	return Join(k1, k2), nil
 }
 
+func (p pairKeyCodec[K1, K2]) Name() string {
+	return fmt.Sprintf("%s,%s", p.key1Name, p.key2Name)
+}
+
+func (p pairKeyCodec[K1, K2]) SchemaCodec() (codec.SchemaCodec[Pair[K1, K2]], error) {
+	field1, err := getNamedKeyField(p.keyCodec1, p.key1Name)
+	if err != nil {
+		return codec.SchemaCodec[Pair[K1, K2]]{}, fmt.Errorf("error getting key1 field: %w", err)
+	}
+
+	field2, err := getNamedKeyField(p.keyCodec2, p.key2Name)
+	if err != nil {
+		return codec.SchemaCodec[Pair[K1, K2]]{}, fmt.Errorf("error getting key2 field: %w", err)
+	}
+
+	return codec.SchemaCodec[Pair[K1, K2]]{
+		Fields: []schema.Field{field1, field2},
+	}, nil
+}
+
+func getNamedKeyField[T any](keyCdc codec.KeyCodec[T], name string) (schema.Field, error) {
+	keySchema, err := codec.KeySchemaCodec(keyCdc)
+	if err != nil {
+		return schema.Field{}, err
+	}
+	if len(keySchema.Fields) != 1 {
+		return schema.Field{}, fmt.Errorf("key schema in composite key has more than one field, got %v", keySchema.Fields)
+	}
+	field := keySchema.Fields[0]
+	field.Name = name
+	return field, nil
+}
+
 // NewPrefixUntilPairRange defines a collection query which ranges until the provided Pair prefix.
 // Unstable: this API might change in the future.
 func NewPrefixUntilPairRange[K1, K2 any](prefix K1) *PairRange[K1, K2] {

collections/quad.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"strings"
 
+	"cosmossdk.io/schema"
+
 	"cosmossdk.io/collections/codec"
 )
 
@@ -79,11 +81,28 @@ func QuadKeyCodec[K1, K2, K3, K4 any](keyCodec1 codec.KeyCodec[K1], keyCodec2 co
 	}
 }
 
+// NamedQuadKeyCodec instantiates a new KeyCodec instance that can encode the Quad, given
+// the KeyCodecs of the four parts of the key, in order.
+// The provided names are used to identify the parts of the key in the schema for indexing.
+func NamedQuadKeyCodec[K1, K2, K3, K4 any](key1Name string, keyCodec1 codec.KeyCodec[K1], key2Name string, keyCodec2 codec.KeyCodec[K2], key3Name string, keyCodec3 codec.KeyCodec[K3], key4Name string, keyCodec4 codec.KeyCodec[K4]) codec.KeyCodec[Quad[K1, K2, K3, K4]] {
+	return quadKeyCodec[K1, K2, K3, K4]{
+		name1:     key1Name,
+		keyCodec1: keyCodec1,
+		name2:     key2Name,
+		keyCodec2: keyCodec2,
+		name3:     key3Name,
+		keyCodec3: keyCodec3,
+		name4:     key4Name,
+		keyCodec4: keyCodec4,
+	}
+}
+
 type quadKeyCodec[K1, K2, K3, K4 any] struct {
-	keyCodec1 codec.KeyCodec[K1]
-	keyCodec2 codec.KeyCodec[K2]
-	keyCodec3 codec.KeyCodec[K3]
-	keyCodec4 codec.KeyCodec[K4]
+	name1, name2, name3, name4 string
+	keyCodec1                  codec.KeyCodec[K1]
+	keyCodec2                  codec.KeyCodec[K2]
+	keyCodec3                  codec.KeyCodec[K3]
+	keyCodec4                  codec.KeyCodec[K4]
 }
 
 type jsonQuadKey [4]json.RawMessage
@@ -338,6 +357,32 @@ func (t quadKeyCodec[K1, K2, K3, K4]) SizeNonTerminal(key Quad[K1, K2, K3, K4])
 	return size
 }
 
+func (t quadKeyCodec[K1, K2, K3, K4]) SchemaCodec() (codec.SchemaCodec[Quad[K1, K2, K3, K4]], error) {
+	field1, err := getNamedKeyField(t.keyCodec1, t.name1)
+	if err != nil {
+		return codec.SchemaCodec[Quad[K1, K2, K3, K4]]{}, fmt.Errorf("error getting key1 field: %w", err)
+	}
+
+	field2, err := getNamedKeyField(t.keyCodec2, t.name2)
+	if err != nil {
+		return codec.SchemaCodec[Quad[K1, K2, K3, K4]]{}, fmt.Errorf("error getting key2 field: %w", err)
+	}
+
+	field3, err := getNamedKeyField(t.keyCodec3, t.name3)
+	if err != nil {
+		return codec.SchemaCodec[Quad[K1, K2, K3, K4]]{}, fmt.Errorf("error getting key3 field: %w", err)
+	}
+
+	field4, err := getNamedKeyField(t.keyCodec4, t.name4)
+	if err != nil {
+		return codec.SchemaCodec[Quad[K1, K2, K3, K4]]{}, fmt.Errorf("error getting key4 field: %w", err)
+	}
+
+	return codec.SchemaCodec[Quad[K1, K2, K3, K4]]{
+		Fields: []schema.Field{field1, field2, field3, field4},
+	}, nil
+}
+
 // NewPrefixUntilQuadRange defines a collection query which ranges until the provided Quad prefix.
 // Unstable: this API might change in the future.
 func NewPrefixUntilQuadRange[K1, K2, K3, K4 any](k1 K1) Ranger[Quad[K1, K2, K3, K4]] {

collections/triple.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"strings"
 
+	"cosmossdk.io/schema"
+
 	"cosmossdk.io/collections/codec"
 )
 
@@ -64,10 +66,22 @@ func TripleKeyCodec[K1, K2, K3 any](keyCodec1 codec.KeyCodec[K1], keyCodec2 code
 	}
 }
 
+func NamedTripleKeyCodec[K1, K2, K3 any](key1Name string, keyCodec1 codec.KeyCodec[K1], key2Name string, keyCodec2 codec.KeyCodec[K2], key3Name string, keyCodec3 codec.KeyCodec[K3]) codec.KeyCodec[Triple[K1, K2, K3]] {
+	return tripleKeyCodec[K1, K2, K3]{
+		key1Name:  key1Name,
+		key2Name:  key2Name,
+		key3Name:  key3Name,
+		keyCodec1: keyCodec1,
+		keyCodec2: keyCodec2,
+		keyCodec3: keyCodec3,
+	}
+}
+
 type tripleKeyCodec[K1, K2, K3 any] struct {
-	keyCodec1 codec.KeyCodec[K1]
-	keyCodec2 codec.KeyCodec[K2]
-	keyCodec3 codec.KeyCodec[K3]
+	key1Name, key2Name, key3Name string
+	keyCodec1                    codec.KeyCodec[K1]
+	keyCodec2                    codec.KeyCodec[K2]
+	keyCodec3                    codec.KeyCodec[K3]
 }
 
 type jsonTripleKey [3]json.RawMessage
@@ -273,6 +287,31 @@ func (t tripleKeyCodec[K1, K2, K3]) SizeNonTerminal(key Triple[K1, K2, K3]) int
 	return size
 }
 
+func (t tripleKeyCodec[K1, K2, K3]) Name() string {
+	return fmt.Sprintf("%s,%s,%s", t.key1Name, t.key2Name, t.key3Name)
+}
+
+func (t tripleKeyCodec[K1, K2, K3]) SchemaCodec() (codec.SchemaCodec[Triple[K1, K2, K3]], error) {
+	field1, err := getNamedKeyField(t.keyCodec1, t.key1Name)
+	if err != nil {
+		return codec.SchemaCodec[Triple[K1, K2, K3]]{}, fmt.Errorf("error getting key1 field: %w", err)
+	}
+
+	field2, err := getNamedKeyField(t.keyCodec2, t.key2Name)
+	if err != nil {
+		return codec.SchemaCodec[Triple[K1, K2, K3]]{}, fmt.Errorf("error getting key2 field: %w", err)
+	}
+
+	field3, err := getNamedKeyField(t.keyCodec3, t.key3Name)
+	if err != nil {
+		return codec.SchemaCodec[Triple[K1, K2, K3]]{}, fmt.Errorf("error getting key3 field: %w", err)
+	}
+
+	return codec.SchemaCodec[Triple[K1, K2, K3]]{
+		Fields: []schema.Field{field1, field2, field3},
+	}, nil
+}
+
 // NewPrefixUntilTripleRange defines a collection query which ranges until the provided Pair prefix.
 // Unstable: this API might change in the future.
 func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1) Ranger[Triple[K1, K2, K3]] {