feat(collections): Infer Indexes in IndexedMap using reflection (#19343)

collections/CHANGELOG.md

@@ -33,7 +33,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Features
 
-* [#18933](https://github.com/cosmos/cosmos-sdk/pull/18933) Add  LookupMap implementation. It is basic wrapping of the standard Map methods but is not iterable.
+* [#19343](https://github.com/cosmos/cosmos-sdk/pull/19343) – Simplify IndexedMap creation by allowing to infer indexes through reflection.
+* [#18933](https://github.com/cosmos/cosmos-sdk/pull/18933) – Add  LookupMap implementation. It is basic wrapping of the standard Map methods but is not iterable.
 * [#17656](https://github.com/cosmos/cosmos-sdk/pull/17656) – Introduces `Vec`, a collection type that allows to represent a growable array on top of a KVStore.
 
 ## [v0.4.0](https://github.com/cosmos/cosmos-sdk/releases/tag/collections%2Fv0.4.0)

collections/README.md

@@ -837,10 +837,6 @@ type AccountsIndexes struct {
 	Number *indexes.Unique[uint64, sdk.AccAddress, authtypes.BaseAccount]
 }
 
-func (a AccountsIndexes) IndexesList() []collections.Index[sdk.AccAddress, authtypes.BaseAccount] {
-	return []collections.Index[sdk.AccAddress, authtypes.BaseAccount]{a.Number}
-}
-
 func NewAccountIndexes(sb *collections.SchemaBuilder) AccountsIndexes {
 	return AccountsIndexes{
 		Number: indexes.NewUnique(
@@ -867,15 +863,23 @@ Where the first type parameter is `uint64`, which is the field type of our index
 The second type parameter is the primary key `sdk.AccAddress`
 And the third type parameter is the actual object we're storing `authtypes.BaseAccount`.
 
-Then we implement a function called `IndexesList` on our `AccountIndexes` struct, this will be used
-by the `IndexedMap` to keep the underlying map in sync with the indexes, in our case `Number`.
-This function just needs to return the slice of indexes contained in the struct.
-
 Then we create a `NewAccountIndexes` function that instantiates and returns the `AccountsIndexes` struct.
 
 The function takes a `SchemaBuilder`. Then we instantiate our `indexes.Unique`, let's analyse the arguments we pass to
 `indexes.NewUnique`.
 
+#### NOTE: indexes list
+
+The `AccountsIndexes` struct contains the indexes, the `NewIndexedMap` function will infer the indexes form that struct
+using reflection, this happens only at init and is not computationally expensive. In case you want to explicitly declare
+indexes: implement the `Indexes` interface in the `AccountsIndexes` struct:
+
+```go
+func (a AccountsIndexes) IndexesList() []collections.Index[sdk.AccAddress, authtypes.BaseAccount] {
+    return []collections.Index[sdk.AccAddress, authtypes.BaseAccount]{a.Number}
+}
+```
+
 #### Instantiating a `indexes.Unique`
 
 The first three arguments, we already know them, they are: `SchemaBuilder`, `Prefix` which is our index prefix (the partition

collections/indexed_map.go

@@ -2,6 +2,9 @@ package collections
 
 import (
 	"context"
+	"errors"
+	"fmt"
+	"reflect"
 
 	"cosmossdk.io/collections/codec"
 )
@@ -32,17 +35,77 @@ type Index[PrimaryKey, Value any] interface {
 // Internally IndexedMap can be seen as a partitioned collection, one partition
 // is a Map[PrimaryKey, Value], that maintains the object, the second
 // are the Indexes.
-type IndexedMap[PrimaryKey, Value any, Idx Indexes[PrimaryKey, Value]] struct {
-	Indexes Idx
-	m       Map[PrimaryKey, Value]
+type IndexedMap[PrimaryKey, Value, Idx any] struct {
+	Indexes         Idx
+	computedIndexes []Index[PrimaryKey, Value]
+	m               Map[PrimaryKey, Value]
+}
+
+// NewIndexedMapSafe behaves like NewIndexedMap but returns errors.
+func NewIndexedMapSafe[K, V, I any](
+	schema *SchemaBuilder,
+	prefix Prefix,
+	name string,
+	pkCodec codec.KeyCodec[K],
+	valueCodec codec.ValueCodec[V],
+	indexes I,
+) (im *IndexedMap[K, V, I], err error) {
+	var indexesList []Index[K, V]
+	indexesImpl, ok := any(indexes).(Indexes[K, V])
+	if ok {
+		indexesList = indexesImpl.IndexesList()
+	} else {
+		// if does not implement Indexes, then we try to infer using reflection
+		indexesList, err = tryInferIndexes[I, K, V](indexes)
+		if err != nil {
+			return nil, fmt.Errorf("unable to infer indexes using reflection, consider implementing Indexes interface: %w", err)
+		}
+	}
+
+	return &IndexedMap[K, V, I]{
+		computedIndexes: indexesList,
+		Indexes:         indexes,
+		m:               NewMap(schema, prefix, name, pkCodec, valueCodec),
+	}, nil
+}
+
+var (
+	// testing sentinel errors
+	errNotStruct = errors.New("wanted struct or pointer to a struct")
+	errNotIndex  = errors.New("field is not an index implementation")
+)
+
+func tryInferIndexes[I, K, V any](indexes I) ([]Index[K, V], error) {
+	typ := reflect.TypeOf(indexes)
+	v := reflect.ValueOf(indexes)
+	// check if struct or pointer to a struct
+	if typ.Kind() != reflect.Struct && (typ.Kind() != reflect.Pointer || typ.Elem().Kind() != reflect.Struct) {
+		return nil, fmt.Errorf("%w: type %v", errNotStruct, typ)
+	}
+	// dereference
+	if typ.Kind() == reflect.Pointer {
+		v = v.Elem()
+	}
+	indexesImpl := make([]Index[K, V], v.NumField())
+	for i := 0; i < v.NumField(); i++ {
+		field := v.Field(i)
+		index, ok := field.Interface().(Index[K, V])
+		if !ok {
+			return nil, fmt.Errorf("%w: field number %d", errNotIndex, i)
+		}
+		indexesImpl[i] = index
+	}
+	return indexesImpl, nil
 }
 
 // NewIndexedMap instantiates a new IndexedMap. Accepts a SchemaBuilder, a Prefix,
 // a humanized name that defines the name of the collection, the primary key codec
 // which is basically what IndexedMap uses to encode the primary key to bytes,
 // the value codec which is what the IndexedMap uses to encode the value.
-// Then it expects the initialized indexes.
-func NewIndexedMap[PrimaryKey, Value any, Idx Indexes[PrimaryKey, Value]](
+// Then it expects the initialized indexes. Reflection is used to infer the
+// indexes, Indexes can optionally be implemented to be explicit. Panics
+// on failure to create indexes. If you want an erroring API use NewIndexedMapSafe.
+func NewIndexedMap[PrimaryKey, Value, Idx any](
 	schema *SchemaBuilder,
 	prefix Prefix,
 	name string,
@@ -50,10 +113,11 @@ func NewIndexedMap[PrimaryKey, Value any, Idx Indexes[PrimaryKey, Value]](
 	valueCodec codec.ValueCodec[Value],
 	indexes Idx,
 ) *IndexedMap[PrimaryKey, Value, Idx] {
-	return &IndexedMap[PrimaryKey, Value, Idx]{
-		Indexes: indexes,
-		m:       NewMap(schema, prefix, name, pkCodec, valueCodec),
+	im, err := NewIndexedMapSafe(schema, prefix, name, pkCodec, valueCodec, indexes)
+	if err != nil {
+		panic(err)
 	}
+	return im
 }
 
 // Get gets the object given its primary key.
@@ -111,7 +175,7 @@ func (m *IndexedMap[PrimaryKey, Value, Idx]) ValueCodec() codec.ValueCodec[Value
 }
 
 func (m *IndexedMap[PrimaryKey, Value, Idx]) ref(ctx context.Context, pk PrimaryKey, value Value) error {
-	for _, index := range m.Indexes.IndexesList() {
+	for _, index := range m.computedIndexes {
 		err := index.Reference(ctx, pk, value, cachedGet[PrimaryKey, Value](ctx, m, pk))
 		if err != nil {
 			return err
@@ -121,7 +185,7 @@ func (m *IndexedMap[PrimaryKey, Value, Idx]) ref(ctx context.Context, pk Primary
 }
 
 func (m *IndexedMap[PrimaryKey, Value, Idx]) unref(ctx context.Context, pk PrimaryKey) error {
-	for _, index := range m.Indexes.IndexesList() {
+	for _, index := range m.computedIndexes {
 		err := index.Unreference(ctx, pk, cachedGet[PrimaryKey, Value](ctx, m, pk))
 		if err != nil {
 			return err

collections/indexed_map_internal_test.go

@@ -0,0 +1,30 @@
+package collections
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTryInferIndex(t *testing.T) {
+	invalidIdx := 5
+
+	t.Run("not a pointer to struct", func(t *testing.T) {
+		_, err := tryInferIndexes[*int, string, string](&invalidIdx)
+		require.ErrorIs(t, err, errNotStruct)
+	})
+
+	t.Run("not a struct", func(t *testing.T) {
+		_, err := tryInferIndexes[int, string, string](invalidIdx)
+		require.ErrorIs(t, err, errNotStruct)
+	})
+
+	t.Run("not an index field", func(t *testing.T) {
+		type invalidIndex struct {
+			A int
+		}
+
+		_, err := tryInferIndexes[invalidIndex, string, string](invalidIndex{})
+		require.ErrorIs(t, err, errNotIndex)
+	})
+}

collections/indexed_map_test.go

@@ -104,3 +104,27 @@ func TestIndexedMap(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, company{"milan", 4}, v)
 }
+
+type inferIndex struct {
+	City *indexes.Multi[string, string, company]
+	Vat  *indexes.Unique[uint64, string, company]
+}
+
+func newInferIndex(schema *collections.SchemaBuilder) *inferIndex {
+	return &inferIndex{
+		City: indexes.NewMulti(schema, collections.NewPrefix(1), "companies_by_city", collections.StringKey, collections.StringKey, func(pk string, value company) (string, error) {
+			return value.City, nil
+		}),
+		Vat: indexes.NewUnique(schema, collections.NewPrefix(2), "companies_by_vat", collections.Uint64Key, collections.StringKey, func(pk string, value company) (uint64, error) {
+			return value.Vat, nil
+		}),
+	}
+}
+
+func TestIndexedMapInfer(t *testing.T) {
+	sk, _ := colltest.MockStore()
+	schema := collections.NewSchemaBuilder(sk)
+
+	_, err := collections.NewIndexedMapSafe(schema, collections.NewPrefix(0), "im", collections.StringKey, colltest.MockValueCodec[company](), newInferIndex(schema))
+	require.NoError(t, err)
+}