fix(indexer): the issues during simapp v1 integration (#22413)

(cherry picked from commit 2290c5e)

# Conflicts:
#	collections/codec/indexing.go
#	collections/indexes/multi.go
#	collections/indexes/reverse_pair.go
#	collections/indexing.go
#	collections/keyset.go
#	collections/map.go
#	collections/pair.go
#	collections/quad.go
#	collections/triple.go
#	simapp/go.mod
#	tests/go.mod

collections/codec/indexing.go

@@ -0,0 +1,102 @@
+package codec
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"cosmossdk.io/schema"
+)
+
+// HasSchemaCodec is an interface that all codec's should implement in order
+// to properly support indexing. It is not required by KeyCodec or ValueCodec
+// in order to preserve backwards compatibility, but a future version of collections
+// may make it required and all codec's should aim to implement it. If it is not
+// implemented, fallback defaults will be used for indexing that may be sub-optimal.
+//
+// Implementations of HasSchemaCodec should test that they are conformant using
+// schema.ValidateObjectKey or schema.ValidateObjectValue depending on whether
+// the codec is a KeyCodec or ValueCodec respectively.
+type HasSchemaCodec[T any] interface {
+	// SchemaCodec returns the schema codec for the collections codec.
+	SchemaCodec() (SchemaCodec[T], error)
+}
+
+// SchemaCodec is a codec that supports converting collection codec values to and
+// from schema codec values.
+type SchemaCodec[T any] struct {
+	// Fields are the schema fields that the codec represents. If this is empty,
+	// it will be assumed that this codec represents no value (such as an item key
+	// or key set value).
+	Fields []schema.Field
+
+	// ToSchemaType converts a codec value of type T to a value corresponding to
+	// a schema object key or value (depending on whether this is a key or value
+	// codec). The returned value should pass validation with schema.ValidateObjectKey
+	// or schema.ValidateObjectValue with the fields specified in Fields.
+	// If this function is nil, it will be assumed that T already represents a
+	// value that conforms to a schema value without any further conversion.
+	ToSchemaType func(T) (any, error)
+
+	// FromSchemaType converts a schema object key or value to T.
+	// If this function is nil, it will be assumed that T already represents a
+	// value that conforms to a schema value without any further conversion.
+	FromSchemaType func(any) (T, error)
+}
+
+// KeySchemaCodec gets the schema codec for the provided KeyCodec either
+// by casting to HasSchemaCodec or returning a fallback codec.
+func KeySchemaCodec[K any](cdc KeyCodec[K]) (SchemaCodec[K], error) {
+	if indexable, ok := cdc.(HasSchemaCodec[K]); ok {
+		return indexable.SchemaCodec()
+	} else {
+		return FallbackSchemaCodec[K](), nil
+	}
+}
+
+// ValueSchemaCodec gets the schema codec for the provided ValueCodec either
+// by casting to HasSchemaCodec or returning a fallback codec.
+func ValueSchemaCodec[V any](cdc ValueCodec[V]) (SchemaCodec[V], error) {
+	if indexable, ok := cdc.(HasSchemaCodec[V]); ok {
+		return indexable.SchemaCodec()
+	} else {
+		return FallbackSchemaCodec[V](), nil
+	}
+}
+
+// FallbackSchemaCodec returns a fallback schema codec for T when one isn't explicitly
+// specified with HasSchemaCodec. It maps all simple types directly to schema kinds
+// and converts everything else to JSON String.
+func FallbackSchemaCodec[T any]() SchemaCodec[T] {
+	var t T
+	kind := schema.KindForGoValue(t)
+	if err := kind.Validate(); err == nil {
+		return SchemaCodec[T]{
+			Fields: []schema.Field{{
+				// we don't set any name so that this can be set to a good default by the caller
+				Name: "",
+				Kind: kind,
+			}},
+			// these can be nil because T maps directly to a schema value for this kind
+			ToSchemaType:   nil,
+			FromSchemaType: nil,
+		}
+	} else {
+		// we default to encoding everything to JSON String
+		return SchemaCodec[T]{
+			Fields: []schema.Field{{Kind: schema.StringKind}},
+			ToSchemaType: func(t T) (any, error) {
+				bz, err := json.Marshal(t)
+				return string(json.RawMessage(bz)), err
+			},
+			FromSchemaType: func(a any) (T, error) {
+				var t T
+				sz, ok := a.(string)
+				if !ok {
+					return t, fmt.Errorf("expected string, got %T", a)
+				}
+				err := json.Unmarshal([]byte(sz), &t)
+				return t, err
+			},
+		}
+	}
+}

collections/indexes/multi.go

@@ -0,0 +1,186 @@
+package indexes
+
+import (
+	"context"
+	"errors"
+
+	"cosmossdk.io/collections"
+	"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.
+type Multi[ReferenceKey, PrimaryKey, Value any] struct {
+	getRefKey func(pk PrimaryKey, value Value) (ReferenceKey, error)
+	refKeys   collections.KeySet[collections.Pair[ReferenceKey, PrimaryKey]]
+}
+
+// NewMulti instantiates a new Multi instance given a schema,
+// a Prefix, the humanized name for the index, the reference key key codec
+// and the primary key key codec. The getRefKeyFunc is a function that
+// given the primary key and value returns the referencing key.
+func NewMulti[ReferenceKey, PrimaryKey, Value any](
+	schema *collections.SchemaBuilder,
+	prefix collections.Prefix,
+	name string,
+	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(),
+				collections.WithKeySetSecondaryIndex(),
+			),
+		}
+	}
+
+	return &Multi[ReferenceKey, PrimaryKey, Value]{
+		getRefKey: getRefKeyFunc,
+		refKeys: collections.NewKeySet(
+			schema,
+			prefix,
+			name,
+			collections.PairKeyCodec(refCodec, pkCodec),
+			collections.WithKeySetSecondaryIndex(),
+		),
+	}
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) Reference(ctx context.Context, pk PrimaryKey, newValue Value, lazyOldValue func() (Value, error)) error {
+	oldValue, err := lazyOldValue()
+	switch {
+	// if no error it means the value existed, and we need to remove the old indexes
+	case err == nil:
+		err = m.unreference(ctx, pk, oldValue)
+		if err != nil {
+			return err
+		}
+	// if error is ErrNotFound, it means that the object does not exist, so we're creating indexes for the first time.
+	// we do nothing.
+	case errors.Is(err, collections.ErrNotFound):
+	// default case means that there was some other error
+	default:
+		return err
+	}
+	// create new indexes
+	refKey, err := m.getRefKey(pk, newValue)
+	if err != nil {
+		return err
+	}
+	return m.refKeys.Set(ctx, collections.Join(refKey, pk))
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) Unreference(ctx context.Context, pk PrimaryKey, getValue func() (Value, error)) error {
+	value, err := getValue()
+	if err != nil {
+		return err
+	}
+	return m.unreference(ctx, pk, value)
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) unreference(ctx context.Context, pk PrimaryKey, value Value) error {
+	refKey, err := m.getRefKey(pk, value)
+	if err != nil {
+		return err
+	}
+	return m.refKeys.Remove(ctx, collections.Join(refKey, pk))
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) Iterate(ctx context.Context, ranger collections.Ranger[collections.Pair[ReferenceKey, PrimaryKey]]) (MultiIterator[ReferenceKey, PrimaryKey], error) {
+	iter, err := m.refKeys.Iterate(ctx, ranger)
+	return (MultiIterator[ReferenceKey, PrimaryKey])(iter), err
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) Walk(
+	ctx context.Context,
+	ranger collections.Ranger[collections.Pair[ReferenceKey, PrimaryKey]],
+	walkFunc func(indexingKey ReferenceKey, indexedKey PrimaryKey) (stop bool, err error),
+) error {
+	return m.refKeys.Walk(ctx, ranger, func(key collections.Pair[ReferenceKey, PrimaryKey]) (bool, error) {
+		return walkFunc(key.K1(), key.K2())
+	})
+}
+
+// MatchExact returns a MultiIterator containing all the primary keys referenced by the provided reference key.
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) MatchExact(ctx context.Context, refKey ReferenceKey) (MultiIterator[ReferenceKey, PrimaryKey], error) {
+	return m.Iterate(ctx, collections.NewPrefixedPairRange[ReferenceKey, PrimaryKey](refKey))
+}
+
+func (m *Multi[K1, K2, Value]) KeyCodec() codec.KeyCodec[collections.Pair[K1, K2]] {
+	return m.refKeys.KeyCodec()
+}
+
+// MultiIterator is just a KeySetIterator with key as Pair[ReferenceKey, PrimaryKey].
+type MultiIterator[ReferenceKey, PrimaryKey any] collections.KeySetIterator[collections.Pair[ReferenceKey, PrimaryKey]]
+
+// PrimaryKey returns the iterator's current primary key.
+func (i MultiIterator[ReferenceKey, PrimaryKey]) PrimaryKey() (PrimaryKey, error) {
+	fullKey, err := i.FullKey()
+	return fullKey.K2(), err
+}
+
+// PrimaryKeys fully consumes the iterator and returns the list of primary keys.
+func (i MultiIterator[ReferenceKey, PrimaryKey]) PrimaryKeys() ([]PrimaryKey, error) {
+	fullKeys, err := i.FullKeys()
+	if err != nil {
+		return nil, err
+	}
+	pks := make([]PrimaryKey, len(fullKeys))
+	for i, fullKey := range fullKeys {
+		pks[i] = fullKey.K2()
+	}
+	return pks, nil
+}
+
+// FullKey returns the current full reference key as Pair[ReferenceKey, PrimaryKey].
+func (i MultiIterator[ReferenceKey, PrimaryKey]) FullKey() (collections.Pair[ReferenceKey, PrimaryKey], error) {
+	return (collections.KeySetIterator[collections.Pair[ReferenceKey, PrimaryKey]])(i).Key()
+}
+
+// FullKeys fully consumes the iterator and returns all the list of full reference keys.
+func (i MultiIterator[ReferenceKey, PrimaryKey]) FullKeys() ([]collections.Pair[ReferenceKey, PrimaryKey], error) {
+	return (collections.KeySetIterator[collections.Pair[ReferenceKey, PrimaryKey]])(i).Keys()
+}
+
+// Next advances the iterator.
+func (i MultiIterator[ReferenceKey, PrimaryKey]) Next() {
+	(collections.KeySetIterator[collections.Pair[ReferenceKey, PrimaryKey]])(i).Next()
+}
+
+// Valid asserts if the iterator is still valid or not.
+func (i MultiIterator[ReferenceKey, PrimaryKey]) Valid() bool {
+	return (collections.KeySetIterator[collections.Pair[ReferenceKey, PrimaryKey]])(i).Valid()
+}
+
+// Close closes the iterator.
+func (i MultiIterator[ReferenceKey, PrimaryKey]) Close() error {
+	return (collections.KeySetIterator[collections.Pair[ReferenceKey, PrimaryKey]])(i).Close()
+}