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

Author: aaronc

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 (