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

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] {