feat(collections): Multi RefKeys method and reverse Triple iterator support

collections/indexes/multi.go

@@ -3,6 +3,8 @@
 import (
 	"context"
 	"errors"
+	"fmt"
+	"reflect"
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/collections/codec"
@@ -121,6 +123,40 @@
 	return m.Iterate(ctx, collections.NewPrefixedPairRange[ReferenceKey, PrimaryKey](refKey))
 }
 
+// RefKeys returns a list of all the MultiIterator's reference keys (may contain duplicates).
+// Enable the "unique" argument to get a unique list of reference keys (the reference key must be comparable)
+func (m *Multi[ReferenceKey, PrimaryKey, Value]) RefKeys(ctx context.Context, unique bool) ([]ReferenceKey, error) {
+	// sanity check - enabled unique with non-comparable ReferenceKey type
+	if unique && !reflect.ValueOf((*ReferenceKey)(nil)).Comparable() {
+		return nil, fmt.Errorf("cannot retrieve unique reference keys since type is not comparable: %T", reflect.TypeOf((*ReferenceKey)(nil)))
+	}
+
+	iter, err := m.refKeys.Iterate(ctx, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	keys := []ReferenceKey{}
+	visited := map[interface{}]struct{}{}
+	for ; iter.Valid(); iter.Next() {
+		key, err := iter.Key()
+		if err != nil {
+			return nil, err
+		}
+		refKey := key.K1()
+
+		if unique {
+			if _, ok := visited[refKey]; ok {
+				continue
+			}
+			visited[refKey] = struct{}{}
+		}
+		keys = append(keys, key.K1())
+	}
+
+	return keys, nil
+}
+
 func (m *Multi[K1, K2, Value]) KeyCodec() codec.KeyCodec[collections.Pair[K1, K2]] {
 	return m.refKeys.KeyCodec()
 }

collections/indexes/multi_test.go

@@ -26,6 +26,14 @@ func TestMultiIndex(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, []uint64{1, 2}, pks)
 
+	// we get all reference keys, should only be "milan"
+	rks, err := mi.RefKeys(ctx, false)
+	require.NoError(t, err)
+	require.Equal(t, []string{"milan", "milan"}, rks)
+	rks, err = mi.RefKeys(ctx, true)
+	require.NoError(t, err)
+	require.Equal(t, []string{"milan"}, rks)
+
 	// replace
 	require.NoError(t, mi.Reference(ctx, 1, company{City: "new york"}, func() (company, error) { return company{City: "milan"}, nil }))
 
@@ -43,6 +51,11 @@ func TestMultiIndex(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, []uint64{1}, pks)
 
+	// assert after replace the reference keys should be "milan" and "new york"
+	rks, err = mi.RefKeys(ctx, false)
+	require.NoError(t, err)
+	require.ElementsMatch(t, []string{"milan", "new york"}, rks)
+
 	// test iter methods
 	iter, err = mi.Iterate(ctx, nil)
 	require.NoError(t, err)

collections/triple.go

@@ -339,3 +339,36 @@ func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2) Ranger[Triple[K1,
 		end:   RangeKeyPrefixEnd(key),
 	}
 }
+
+// NewPrefixUntilTripleRangeReversed defines a collection query which ranges until the provided Pair prefix
+// in reverse order.
+// Unstable: this API might change in the future.
+func NewPrefixUntilTripleRangeReversed[K1, K2, K3 any](k1 K1) Ranger[Triple[K1, K2, K3]] {
+	key := TriplePrefix[K1, K2, K3](k1)
+	return &Range[Triple[K1, K2, K3]]{
+		end:   RangeKeyPrefixEnd(key),
+		order: OrderDescending,
+	}
+}
+
+// NewPrefixedTripleRange provides a Range for all keys prefixed with the given
+// first part of the Triple key in reverse order.
+func NewPrefixedTripleRangeReversed[K1, K2, K3 any](k1 K1) Ranger[Triple[K1, K2, K3]] {
+	key := TriplePrefix[K1, K2, K3](k1)
+	return &Range[Triple[K1, K2, K3]]{
+		start: RangeKeyExact(key),
+		end:   RangeKeyPrefixEnd(key),
+		order: OrderDescending,
+	}
+}
+
+// NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given
+// first and second parts of the Triple key in reverse order.
+func NewSuperPrefixedTripleRangeReversed[K1, K2, K3 any](k1 K1, k2 K2) Ranger[Triple[K1, K2, K3]] {
+	key := TripleSuperPrefix[K1, K2, K3](k1, k2)
+	return &Range[Triple[K1, K2, K3]]{
+		start: RangeKeyExact(key),
+		end:   RangeKeyPrefixEnd(key),
+		order: OrderDescending,
+	}
+}

collections/triple_test.go

@@ -7,7 +7,7 @@ import (
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/collections/colltest"
-	"cosmossdk.io/core/testing"
+	coretesting "cosmossdk.io/core/testing"
 )
 
 func TestTriple(t *testing.T) {
@@ -45,10 +45,30 @@ func TestTripleRange(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, keys[:3], gotKeys)
 
+	// we prefix over (1) with "reverse" enabled, we expect 3 results in reverse order
+	iter, err = keySet.Iterate(ctx, collections.NewPrefixedTripleRangeReversed[uint64, string, []byte](uint64(1)))
+	require.NoError(t, err)
+	gotKeys, err = iter.Keys()
+	require.NoError(t, err)
+	require.Len(t, gotKeys, 3)
+	for i := range gotKeys {
+		require.Equal(t, gotKeys[i], keys[len(gotKeys)-i-1])
+	}
+
 	// we super prefix over Join(1, "A") we expect 2 results
 	iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A"))
 	require.NoError(t, err)
 	gotKeys, err = iter.Keys()
 	require.NoError(t, err)
 	require.Equal(t, keys[:2], gotKeys)
+
+	// we prefix over Join(1, "A") with "reverse" enabled, we expect 2 results in reverse order
+	iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRangeReversed[uint64, string, []byte](1, "A"))
+	require.NoError(t, err)
+	gotKeys, err = iter.Keys()
+	require.NoError(t, err)
+	require.Len(t, gotKeys, 2)
+	for i := range gotKeys {
+		require.Equal(t, gotKeys[i], keys[len(gotKeys)-i-1])
+	}
 }