diff --git a/collections/indexes/multi.go b/collections/indexes/multi.go index 49ad2fb6f378..144753b782f0 100644 --- a/collections/indexes/multi.go +++ b/collections/indexes/multi.go @@ -121,6 +121,48 @@ func (m *Multi[ReferenceKey, PrimaryKey, Value]) MatchExact(ctx context.Context, 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) +// WARNING: The use of RefKeys() can be very expensive in terms of Gas. Please make sure you iterate over a relatively +// small set of reference keys. +func (m *Multi[ReferenceKey, PrimaryKey, Value]) RefKeys(ctx context.Context, unique bool) ([]ReferenceKey, error) { + iter, err := m.refKeys.IterateRaw(ctx, nil, nil, collections.OrderAscending) + if err != nil { + return nil, err + } + + keys := []ReferenceKey{} + visited := map[string]struct{}{} + for ; iter.Valid(); iter.Next() { + key, err := iter.Key() + if err != nil { + return nil, err + } + refKey := key.K1() + + if unique { + // encode the ref key using its codec. casting to pairKeyCodec must + // work since by definition the Multi key codec is a pair key codec + // of [ReferenceKey, PrimaryKey] + refKeyCodec := m.refKeys.KeyCodec().(pairKeyCodec[ReferenceKey, PrimaryKey]).KeyCodec1() + buf := make([]byte, refKeyCodec.Size(refKey)) + _, err := refKeyCodec.Encode(buf, refKey) + if err != nil { + return nil, err + } + + // check if visited + if _, ok := visited[string(buf)]; ok { + continue + } + visited[string(buf)] = struct{}{} + } + keys = append(keys, refKey) + } + + return keys, nil +} + func (m *Multi[K1, K2, Value]) KeyCodec() codec.KeyCodec[collections.Pair[K1, K2]] { return m.refKeys.KeyCodec() } diff --git a/collections/indexes/multi_test.go b/collections/indexes/multi_test.go index a0b1313b8568..5fe106cb2916 100644 --- a/collections/indexes/multi_test.go +++ b/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) diff --git a/collections/triple.go b/collections/triple.go index e4a07970945c..f83412059d8c 100644 --- a/collections/triple.go +++ b/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, + } +} + +// NewPrefixedTripleRangeReversed 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, + } +} + +// NewSuperPrefixedTripleRangeReversed 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, + } +} diff --git a/collections/triple_test.go b/collections/triple_test.go index 5ea53b3fe667..c605a7aeaba7 100644 --- a/collections/triple_test.go +++ b/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]) + } }