feat(collections): indexes for IndexedMap (cosmos#14706)

Co-authored-by: testinginprod <testinginprod@somewhere.idk>
Co-authored-by: Marko <marbar3778@yahoo.com>
Co-authored-by: Likhita Polavarapu <78951027+likhita-809@users.noreply.github.com>

collections/indexes/doc.go

@@ -0,0 +1,3 @@
+// Package indexes contains the most common indexes types to be used with a collections.IndexedMap.
+// It also contains specialised helper functions to collect and query efficiently an index.
+package indexes

collections/indexes/helpers.go

@@ -0,0 +1,110 @@
+package indexes
+
+import (
+	"context"
+	"cosmossdk.io/collections"
+)
+
+// Iterator defines the minimum set of methods of an index iterator
+// required to work with the helpers.
+type Iterator[K any] interface {
+	// PrimaryKey returns the iterator current primary key.
+	PrimaryKey() (K, error)
+	// Next advances the iterator by one element.
+	Next()
+	// Valid asserts if the Iterator is valid.
+	Valid() bool
+	// Close closes the iterator.
+	Close() error
+}
+
+// CollectKeyValues collects all the keys and the values of an indexed map index iterator.
+// The Iterator is fully consumed and closed.
+func CollectKeyValues[K, V any, I Iterator[K], Idx collections.Indexes[K, V]](
+	ctx context.Context,
+	indexedMap *collections.IndexedMap[K, V, Idx],
+	iter I) (kvs []collections.KeyValue[K, V], err error) {
+	err = ScanKeyValues(ctx, indexedMap, iter, func(kv collections.KeyValue[K, V]) bool {
+		kvs = append(kvs, kv)
+		return false
+	})
+	return
+}
+
+// ScanKeyValues calls the do function on every record found, in the indexed map
+// from the index iterator. Returning false stops the iteration.
+// The Iterator is closed when this function exits.
+func ScanKeyValues[K, V any, I Iterator[K], Idx collections.Indexes[K, V]](
+	ctx context.Context,
+	indexedMap *collections.IndexedMap[K, V, Idx],
+	iter I,
+	do func(kv collections.KeyValue[K, V]) (stop bool)) (err error) {
+
+	defer iter.Close()
+
+	for ; iter.Valid(); iter.Next() {
+		pk, err := iter.PrimaryKey()
+		if err != nil {
+			return err
+		}
+
+		value, err := indexedMap.Get(ctx, pk)
+		if err != nil {
+			return err
+		}
+
+		kv := collections.KeyValue[K, V]{
+			Key:   pk,
+			Value: value,
+		}
+
+		if do(kv) {
+			break
+		}
+	}
+
+	return nil
+}
+
+// CollectValues collects all the values from an Index iterator and the IndexedMap.
+// Closes the Iterator.
+func CollectValues[K, V any, I Iterator[K], Idx collections.Indexes[K, V]](
+	ctx context.Context,
+	indexedMap *collections.IndexedMap[K, V, Idx],
+	iter I) (values []V, err error) {
+	err = ScanValues(ctx, indexedMap, iter, func(value V) (stop bool) {
+		values = append(values, value)
+		return false
+	})
+	return
+}
+
+// ScanValues collects all the values from an Index iterator and the IndexedMap in a lazy way.
+// The iterator is closed when this function exits.
+func ScanValues[K, V any, I Iterator[K], Idx collections.Indexes[K, V]](
+	ctx context.Context,
+	indexedMap *collections.IndexedMap[K, V, Idx],
+	iter I,
+	f func(value V) (stop bool),
+) error {
+	defer iter.Close()
+
+	for ; iter.Valid(); iter.Next() {
+		key, err := iter.PrimaryKey()
+		if err != nil {
+			return err
+		}
+
+		value, err := indexedMap.Get(ctx, key)
+		if err != nil {
+			return err
+		}
+
+		stop := f(value)
+		if stop {
+			return nil
+		}
+	}
+
+	return nil
+}

collections/indexes/helpers_test.go

@@ -0,0 +1,87 @@
+package indexes
+
+import (
+	"cosmossdk.io/collections"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestHelpers(t *testing.T) {
+	// uses MultiPair scenario.
+	// We store balances as:
+	// Key: Pair[Address=string, Denom=string] => Value: Amount=uint64
+
+	sk, ctx := deps()
+	sb := collections.NewSchemaBuilder(sk)
+
+	keyCodec := collections.PairKeyCodec(collections.StringKey, collections.StringKey)
+	indexedMap := collections.NewIndexedMap(
+		sb,
+		collections.NewPrefix("balances"), "balances",
+		keyCodec,
+		collections.Uint64Value,
+		balanceIndex{
+			Denom: NewMultiPair[Amount](sb, collections.NewPrefix("denom_index"), "denom_index", keyCodec),
+		},
+	)
+
+	err := indexedMap.Set(ctx, collections.Join("address1", "atom"), 100)
+	require.NoError(t, err)
+
+	err = indexedMap.Set(ctx, collections.Join("address1", "osmo"), 200)
+	require.NoError(t, err)
+
+	err = indexedMap.Set(ctx, collections.Join("address2", "osmo"), 300)
+	require.NoError(t, err)
+
+	// test collect values
+	iter, err := indexedMap.Indexes.Denom.MatchExact(ctx, "osmo")
+	require.NoError(t, err)
+
+	values, err := CollectValues(ctx, indexedMap, iter)
+	require.NoError(t, err)
+	require.Equal(t, []Amount{200, 300}, values)
+
+	// test collect key values
+
+	iter, err = indexedMap.Indexes.Denom.MatchExact(ctx, "osmo")
+	require.NoError(t, err)
+	kvs, err := CollectKeyValues(ctx, indexedMap, iter)
+	require.NoError(t, err)
+
+	require.Equal(t, []collections.KeyValue[collections.Pair[Address, Denom], Amount]{
+		{
+			Key:   collections.Join("address1", "osmo"),
+			Value: 200,
+		},
+		{
+			Key:   collections.Join("address2", "osmo"),
+			Value: 300,
+		},
+	}, kvs)
+
+	// test scan values with early termination
+	iter, err = indexedMap.Indexes.Denom.MatchExact(ctx, "osmo")
+	require.NoError(t, err)
+	numCalled := 0
+	err = ScanValues(ctx, indexedMap, iter, func(v Amount) bool {
+		require.Equal(t, Amount(200), v)
+		numCalled++
+		require.Equal(t, numCalled, 1)
+		return true // says to stop
+	})
+	require.NoError(t, err)
+
+	// test scan kv with early termination
+	iter, err = indexedMap.Indexes.Denom.MatchExact(ctx, "osmo")
+	require.NoError(t, err)
+	numCalled = 0
+	err = ScanKeyValues(ctx, indexedMap, iter, func(kv collections.KeyValue[collections.Pair[Address, Denom], Amount]) bool {
+		require.Equal(t, Amount(200), kv.Value)
+		require.Equal(t, collections.Join("address1", "osmo"), kv.Key)
+		numCalled++
+		require.Equal(t, numCalled, 1)
+		return true // says to stop
+	})
+	require.NoError(t, err)
+}

collections/indexes/indexes_test.go

@@ -0,0 +1,53 @@
+package indexes
+
+import (
+	"context"
+	"cosmossdk.io/core/store"
+	db "github.com/cosmos/cosmos-db"
+)
+
+// TODO remove this when we add testStore to core/store.
+
+type testStore struct {
+	db db.DB
+}
+
+func (t testStore) OpenKVStore(ctx context.Context) store.KVStore {
+	return t
+}
+
+func (t testStore) Get(key []byte) ([]byte, error) {
+	return t.db.Get(key)
+}
+
+func (t testStore) Has(key []byte) (bool, error) {
+	return t.db.Has(key)
+}
+
+func (t testStore) Set(key, value []byte) error {
+	return t.db.Set(key, value)
+}
+
+func (t testStore) Delete(key []byte) error {
+	return t.db.Delete(key)
+}
+
+func (t testStore) Iterator(start, end []byte) (store.Iterator, error) {
+	return t.db.Iterator(start, end)
+}
+
+func (t testStore) ReverseIterator(start, end []byte) (store.Iterator, error) {
+	return t.db.ReverseIterator(start, end)
+}
+
+var _ store.KVStore = testStore{}
+
+func deps() (store.KVStoreService, context.Context) {
+	kv := db.NewMemDB()
+	return &testStore{kv}, context.Background()
+}
+
+type company struct {
+	City string
+	Vat  uint64
+}

collections/indexes/multi.go

@@ -0,0 +1,104 @@
+package indexes
+
+import (
+	"context"
+	"cosmossdk.io/collections"
+)
+
+// 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] collections.GenericMultiIndex[ReferenceKey, PrimaryKey, PrimaryKey, Value]
+
+// 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 collections.KeyCodec[ReferenceKey],
+	pkCodec collections.KeyCodec[PrimaryKey],
+	getRefKeyFunc func(pk PrimaryKey, value Value) (ReferenceKey, error),
+) *Multi[ReferenceKey, PrimaryKey, Value] {
+	i := collections.NewGenericMultiIndex(
+		schema, prefix, name, refCodec, pkCodec,
+		func(pk PrimaryKey, value Value) ([]collections.IndexReference[ReferenceKey, PrimaryKey], error) {
+			ref, err := getRefKeyFunc(pk, value)
+			if err != nil {
+				return nil, err
+			}
+			return []collections.IndexReference[ReferenceKey, PrimaryKey]{
+				collections.NewIndexReference(ref, pk),
+			}, nil
+		},
+	)
+
+	return (*Multi[ReferenceKey, PrimaryKey, Value])(i)
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) Reference(ctx context.Context, pk PrimaryKey, newValue Value, oldValue *Value) error {
+	return (*collections.GenericMultiIndex[ReferenceKey, PrimaryKey, PrimaryKey, Value])(m).Reference(ctx, pk, newValue, oldValue)
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) Unreference(ctx context.Context, pk PrimaryKey, value Value) error {
+	return (*collections.GenericMultiIndex[ReferenceKey, PrimaryKey, PrimaryKey, Value])(m).Unreference(ctx, pk, value)
+}
+
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) Iterate(ctx context.Context, ranger collections.Ranger[collections.Pair[ReferenceKey, PrimaryKey]]) (MultiIterator[ReferenceKey, PrimaryKey], error) {
+	iter, err := (*collections.GenericMultiIndex[ReferenceKey, PrimaryKey, PrimaryKey, Value])(m).Iterate(ctx, ranger)
+	return (MultiIterator[ReferenceKey, PrimaryKey])(iter), err
+}
+
+// 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))
+}
+
+// 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()
+}