From e241646aff125d83b4926c8eac44e878499a3a49 Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Sun, 1 Sep 2024 17:26:10 +0300
Subject: [PATCH 1/8] add RefKeys() method + unit test

---
 collections/indexes/multi.go      | 36 +++++++++++++++++++++++++++++++
 collections/indexes/multi_test.go | 13 +++++++++++
 2 files changed, 49 insertions(+)

diff --git a/collections/indexes/multi.go b/collections/indexes/multi.go
index 49ad2fb6f378..6a4a7f2d1a22 100644
--- a/collections/indexes/multi.go
+++ b/collections/indexes/multi.go
@@ -3,6 +3,8 @@ package indexes
 import (
 	"context"
 	"errors"
+	"fmt"
+	"reflect"
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/collections/codec"
@@ -121,6 +123,40 @@ 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)
+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()
 }
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)

From 564cc8e9dccdddb2e2d717fc534a9780cab53aca Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Sun, 1 Sep 2024 17:48:44 +0300
Subject: [PATCH 2/8] added support for a reverse triple iterator

---
 collections/triple.go      | 23 +++++++++++++++++++----
 collections/triple_test.go | 26 +++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 7 deletions(-)

diff --git a/collections/triple.go b/collections/triple.go
index 9733d9984099..8bc20abc90c5 100644
--- a/collections/triple.go
+++ b/collections/triple.go
@@ -275,29 +275,44 @@ func (t tripleKeyCodec[K1, K2, K3]) SizeNonTerminal(key Triple[K1, K2, K3]) int
 
 // NewPrefixUntilTripleRange defines a collection query which ranges until the provided Pair prefix.
 // Unstable: this API might change in the future.
-func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1) Ranger[Triple[K1, K2, K3]] {
+func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1, reverse bool) Ranger[Triple[K1, K2, K3]] {
 	key := TriplePrefix[K1, K2, K3](k1)
+	order := OrderAscending
+	if reverse {
+		order = OrderDescending
+	}
 	return &Range[Triple[K1, K2, K3]]{
-		end: RangeKeyPrefixEnd(key),
+		end:   RangeKeyPrefixEnd(key),
+		order: order,
 	}
 }
 
 // NewPrefixedTripleRange provides a Range for all keys prefixed with the given
 // first part of the Triple key.
-func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1) Ranger[Triple[K1, K2, K3]] {
+func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1, reverse bool) Ranger[Triple[K1, K2, K3]] {
 	key := TriplePrefix[K1, K2, K3](k1)
+	order := OrderAscending
+	if reverse {
+		order = OrderDescending
+	}
 	return &Range[Triple[K1, K2, K3]]{
 		start: RangeKeyExact(key),
 		end:   RangeKeyPrefixEnd(key),
+		order: order,
 	}
 }
 
 // NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given
 // first and second parts of the Triple key.
-func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2) Ranger[Triple[K1, K2, K3]] {
+func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2, reverse bool) Ranger[Triple[K1, K2, K3]] {
 	key := TripleSuperPrefix[K1, K2, K3](k1, k2)
+	order := OrderAscending
+	if reverse {
+		order = OrderDescending
+	}
 	return &Range[Triple[K1, K2, K3]]{
 		start: RangeKeyExact(key),
 		end:   RangeKeyPrefixEnd(key),
+		order: order,
 	}
 }
diff --git a/collections/triple_test.go b/collections/triple_test.go
index 5ea53b3fe667..95e756868028 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) {
@@ -39,16 +39,36 @@ func TestTripleRange(t *testing.T) {
 	}
 
 	// we prefix over (1) we expect 3 results
-	iter, err := keySet.Iterate(ctx, collections.NewPrefixedTripleRange[uint64, string, []byte](uint64(1)))
+	iter, err := keySet.Iterate(ctx, collections.NewPrefixedTripleRange[uint64, string, []byte](uint64(1), false))
 	require.NoError(t, err)
 	gotKeys, err := iter.Keys()
 	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.NewPrefixedTripleRange[uint64, string, []byte](uint64(1), true))
+	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"))
+	iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A", false))
 	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.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A", true))
+	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])
+	}
 }

From e1a181102fc0c910af039c1091fc2806eb3d6a00 Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Sun, 1 Sep 2024 18:24:24 +0300
Subject: [PATCH 3/8] fix triple ranger calls

---
 collections/README.md           | 4 ++--
 x/distribution/keeper/hooks.go  | 2 +-
 x/feegrant/keeper/keeper.go     | 2 +-
 x/staking/keeper/delegation.go  | 6 +++---
 x/staking/keeper/query_utils.go | 2 +-
 5 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/collections/README.md b/collections/README.md
index 298f684e6691..fbcabb1a361f 100644
--- a/collections/README.md
+++ b/collections/README.md
@@ -1166,7 +1166,7 @@ func NewKeeper(storeKey *storetypes.KVStoreKey) Keeper {
 // RedelegationsByDelegator iterates over all the redelegations of a given delegator and calls onResult providing
 // each redelegation from source validator towards the destination validator.
 func (k Keeper) RedelegationsByDelegator(ctx context.Context, delegator AccAddress, onResult func(src, dst ValAddress) (stop bool, err error)) error {
- rng := collections.NewPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator)
+ rng := collections.NewPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator, false)
  return k.Redelegations.Walk(ctx, rng, func(key collections.Triple[AccAddress, ValAddress, ValAddress]) (stop bool, err error) {
   return onResult(key.K2(), key.K3())
  })
@@ -1175,7 +1175,7 @@ func (k Keeper) RedelegationsByDelegator(ctx context.Context, delegator AccAddre
 // RedelegationsByDelegatorAndValidator iterates over all the redelegations of a given delegator and its source validator and calls onResult for each
 // destination validator.
 func (k Keeper) RedelegationsByDelegatorAndValidator(ctx context.Context, delegator AccAddress, validator ValAddress, onResult func(dst ValAddress) (stop bool, err error)) error {
- rng := collections.NewSuperPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator, validator)
+ rng := collections.NewSuperPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator, validator, false)
  return k.Redelegations.Walk(ctx, rng, func(key collections.Triple[AccAddress, ValAddress, ValAddress]) (stop bool, err error) {
   return onResult(key.K3())
  })
diff --git a/x/distribution/keeper/hooks.go b/x/distribution/keeper/hooks.go
index 46605690c26f..729f4b88e346 100644
--- a/x/distribution/keeper/hooks.go
+++ b/x/distribution/keeper/hooks.go
@@ -110,7 +110,7 @@ func (h Hooks) AfterValidatorRemoved(ctx context.Context, _ sdk.ConsAddress, val
 	}
 
 	// clear slashes
-	err = h.k.ValidatorSlashEvents.Clear(ctx, collections.NewPrefixedTripleRange[sdk.ValAddress, uint64, uint64](valAddr))
+	err = h.k.ValidatorSlashEvents.Clear(ctx, collections.NewPrefixedTripleRange[sdk.ValAddress, uint64, uint64](valAddr, false))
 	if err != nil {
 		return err
 	}
diff --git a/x/feegrant/keeper/keeper.go b/x/feegrant/keeper/keeper.go
index b59fd7caa0f5..724ef0b56ac2 100644
--- a/x/feegrant/keeper/keeper.go
+++ b/x/feegrant/keeper/keeper.go
@@ -294,7 +294,7 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*feegrant.GenesisState, erro
 // RemoveExpiredAllowances iterates grantsByExpiryQueue and deletes the expired grants.
 func (k Keeper) RemoveExpiredAllowances(ctx context.Context, limit int) error {
 	exp := k.HeaderService.HeaderInfo(ctx).Time
-	rng := collections.NewPrefixUntilTripleRange[time.Time, sdk.AccAddress, sdk.AccAddress](exp)
+	rng := collections.NewPrefixUntilTripleRange[time.Time, sdk.AccAddress, sdk.AccAddress](exp, false)
 	count := 0
 
 	keysToRemove := []collections.Triple[time.Time, sdk.AccAddress, sdk.AccAddress]{}
diff --git a/x/staking/keeper/delegation.go b/x/staking/keeper/delegation.go
index 69ad3843ee9c..18211132265f 100644
--- a/x/staking/keeper/delegation.go
+++ b/x/staking/keeper/delegation.go
@@ -421,7 +421,7 @@ func (k Keeper) GetRedelegations(ctx context.Context, delegator sdk.AccAddress,
 	redelegations = make([]types.Redelegation, maxRetrieve)
 
 	i := 0
-	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator)
+	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator, false)
 	err = k.Redelegations.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], redelegation types.Redelegation) (stop bool, err error) {
 		if i >= int(maxRetrieve) {
 			return true, nil
@@ -441,7 +441,7 @@ func (k Keeper) GetRedelegations(ctx context.Context, delegator sdk.AccAddress,
 // GetRedelegationsFromSrcValidator returns all redelegations from a particular
 // validator.
 func (k Keeper) GetRedelegationsFromSrcValidator(ctx context.Context, valAddr sdk.ValAddress) (reds []types.Redelegation, err error) {
-	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](valAddr)
+	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](valAddr, false)
 	err = k.RedelegationsByValSrc.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], value []byte) (stop bool, err error) {
 		valSrcAddr, delAddr, valDstAddr := key.K1(), key.K2(), key.K3()
 
@@ -462,7 +462,7 @@ func (k Keeper) GetRedelegationsFromSrcValidator(ctx context.Context, valAddr sd
 
 // HasReceivingRedelegation checks if validator is receiving a redelegation.
 func (k Keeper) HasReceivingRedelegation(ctx context.Context, delAddr sdk.AccAddress, valDstAddr sdk.ValAddress) (bool, error) {
-	rng := collections.NewSuperPrefixedTripleRange[[]byte, []byte, []byte](valDstAddr, delAddr)
+	rng := collections.NewSuperPrefixedTripleRange[[]byte, []byte, []byte](valDstAddr, delAddr, false)
 	hasReceivingRedelegation := false
 	err := k.RedelegationsByValDst.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], value []byte) (stop bool, err error) {
 		hasReceivingRedelegation = true
diff --git a/x/staking/keeper/query_utils.go b/x/staking/keeper/query_utils.go
index 625984dbd44a..ebe23b81ebed 100644
--- a/x/staking/keeper/query_utils.go
+++ b/x/staking/keeper/query_utils.go
@@ -107,7 +107,7 @@ func (k Keeper) GetAllRedelegations(
 	dstValFilter := !(dstValAddress.Empty())
 
 	redelegations := []types.Redelegation{}
-	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator)
+	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator, false)
 	err := k.Redelegations.Walk(ctx, rng,
 		func(key collections.Triple[[]byte, []byte, []byte], redelegation types.Redelegation) (stop bool, err error) {
 			valSrcAddr, valDstAddr := key.K2(), key.K3()

From 1b1d1a7652cf18a9b1d458c8c2526a7432e1c6b8 Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Mon, 2 Sep 2024 14:01:19 +0300
Subject: [PATCH 4/8] add reverse iterator as an optional argument to avoid
 breaking the API

---
 collections/README.md           |  4 ++--
 collections/triple.go           | 36 ++++++++++++++++++++++-----------
 collections/triple_test.go      |  4 ++--
 x/distribution/keeper/hooks.go  |  2 +-
 x/feegrant/keeper/keeper.go     |  2 +-
 x/staking/keeper/delegation.go  |  6 +++---
 x/staking/keeper/query_utils.go |  2 +-
 7 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/collections/README.md b/collections/README.md
index fbcabb1a361f..298f684e6691 100644
--- a/collections/README.md
+++ b/collections/README.md
@@ -1166,7 +1166,7 @@ func NewKeeper(storeKey *storetypes.KVStoreKey) Keeper {
 // RedelegationsByDelegator iterates over all the redelegations of a given delegator and calls onResult providing
 // each redelegation from source validator towards the destination validator.
 func (k Keeper) RedelegationsByDelegator(ctx context.Context, delegator AccAddress, onResult func(src, dst ValAddress) (stop bool, err error)) error {
- rng := collections.NewPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator, false)
+ rng := collections.NewPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator)
  return k.Redelegations.Walk(ctx, rng, func(key collections.Triple[AccAddress, ValAddress, ValAddress]) (stop bool, err error) {
   return onResult(key.K2(), key.K3())
  })
@@ -1175,7 +1175,7 @@ func (k Keeper) RedelegationsByDelegator(ctx context.Context, delegator AccAddre
 // RedelegationsByDelegatorAndValidator iterates over all the redelegations of a given delegator and its source validator and calls onResult for each
 // destination validator.
 func (k Keeper) RedelegationsByDelegatorAndValidator(ctx context.Context, delegator AccAddress, validator ValAddress, onResult func(dst ValAddress) (stop bool, err error)) error {
- rng := collections.NewSuperPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator, validator, false)
+ rng := collections.NewSuperPrefixedTripleRange[AccAddress, ValAddress, ValAddress](delegator, validator)
  return k.Redelegations.Walk(ctx, rng, func(key collections.Triple[AccAddress, ValAddress, ValAddress]) (stop bool, err error) {
   return onResult(key.K3())
  })
diff --git a/collections/triple.go b/collections/triple.go
index 8bc20abc90c5..cb4d8aac2a8e 100644
--- a/collections/triple.go
+++ b/collections/triple.go
@@ -274,13 +274,17 @@ func (t tripleKeyCodec[K1, K2, K3]) SizeNonTerminal(key Triple[K1, K2, K3]) int
 }
 
 // NewPrefixUntilTripleRange defines a collection query which ranges until the provided Pair prefix.
-// Unstable: this API might change in the future.
-func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1, reverse bool) Ranger[Triple[K1, K2, K3]] {
+// Unstable: this API might change in the future. Use the first boolean option to get a reverse iterator (reverse=true).
+func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1, opts ...bool) Ranger[Triple[K1, K2, K3]] {
 	key := TriplePrefix[K1, K2, K3](k1)
 	order := OrderAscending
-	if reverse {
-		order = OrderDescending
+	if len(opts) > 0 {
+		reverse := opts[0]
+		if reverse {
+			order = OrderDescending
+		}
 	}
+
 	return &Range[Triple[K1, K2, K3]]{
 		end:   RangeKeyPrefixEnd(key),
 		order: order,
@@ -288,13 +292,17 @@ func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1, reverse bool) Ranger[Tripl
 }
 
 // NewPrefixedTripleRange provides a Range for all keys prefixed with the given
-// first part of the Triple key.
-func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1, reverse bool) Ranger[Triple[K1, K2, K3]] {
+// first part of the Triple key. Use the first boolean option to get a reverse iterator (reverse=true).
+func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1, opts ...bool) Ranger[Triple[K1, K2, K3]] {
 	key := TriplePrefix[K1, K2, K3](k1)
 	order := OrderAscending
-	if reverse {
-		order = OrderDescending
+	if len(opts) > 0 {
+		reverse := opts[0]
+		if reverse {
+			order = OrderDescending
+		}
 	}
+
 	return &Range[Triple[K1, K2, K3]]{
 		start: RangeKeyExact(key),
 		end:   RangeKeyPrefixEnd(key),
@@ -303,13 +311,17 @@ func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1, reverse bool) Ranger[Triple[K
 }
 
 // NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given
-// first and second parts of the Triple key.
-func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2, reverse bool) Ranger[Triple[K1, K2, K3]] {
+// first and second parts of the Triple key. Use the first boolean option to get a reverse iterator (reverse=true).
+func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2, opts ...bool) Ranger[Triple[K1, K2, K3]] {
 	key := TripleSuperPrefix[K1, K2, K3](k1, k2)
 	order := OrderAscending
-	if reverse {
-		order = OrderDescending
+	if len(opts) > 0 {
+		reverse := opts[0]
+		if reverse {
+			order = OrderDescending
+		}
 	}
+
 	return &Range[Triple[K1, K2, K3]]{
 		start: RangeKeyExact(key),
 		end:   RangeKeyPrefixEnd(key),
diff --git a/collections/triple_test.go b/collections/triple_test.go
index 95e756868028..a6dae56abfee 100644
--- a/collections/triple_test.go
+++ b/collections/triple_test.go
@@ -39,7 +39,7 @@ func TestTripleRange(t *testing.T) {
 	}
 
 	// we prefix over (1) we expect 3 results
-	iter, err := keySet.Iterate(ctx, collections.NewPrefixedTripleRange[uint64, string, []byte](uint64(1), false))
+	iter, err := keySet.Iterate(ctx, collections.NewPrefixedTripleRange[uint64, string, []byte](uint64(1)))
 	require.NoError(t, err)
 	gotKeys, err := iter.Keys()
 	require.NoError(t, err)
@@ -56,7 +56,7 @@ func TestTripleRange(t *testing.T) {
 	}
 
 	// we super prefix over Join(1, "A") we expect 2 results
-	iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A", false))
+	iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A"))
 	require.NoError(t, err)
 	gotKeys, err = iter.Keys()
 	require.NoError(t, err)
diff --git a/x/distribution/keeper/hooks.go b/x/distribution/keeper/hooks.go
index 729f4b88e346..46605690c26f 100644
--- a/x/distribution/keeper/hooks.go
+++ b/x/distribution/keeper/hooks.go
@@ -110,7 +110,7 @@ func (h Hooks) AfterValidatorRemoved(ctx context.Context, _ sdk.ConsAddress, val
 	}
 
 	// clear slashes
-	err = h.k.ValidatorSlashEvents.Clear(ctx, collections.NewPrefixedTripleRange[sdk.ValAddress, uint64, uint64](valAddr, false))
+	err = h.k.ValidatorSlashEvents.Clear(ctx, collections.NewPrefixedTripleRange[sdk.ValAddress, uint64, uint64](valAddr))
 	if err != nil {
 		return err
 	}
diff --git a/x/feegrant/keeper/keeper.go b/x/feegrant/keeper/keeper.go
index 724ef0b56ac2..b59fd7caa0f5 100644
--- a/x/feegrant/keeper/keeper.go
+++ b/x/feegrant/keeper/keeper.go
@@ -294,7 +294,7 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*feegrant.GenesisState, erro
 // RemoveExpiredAllowances iterates grantsByExpiryQueue and deletes the expired grants.
 func (k Keeper) RemoveExpiredAllowances(ctx context.Context, limit int) error {
 	exp := k.HeaderService.HeaderInfo(ctx).Time
-	rng := collections.NewPrefixUntilTripleRange[time.Time, sdk.AccAddress, sdk.AccAddress](exp, false)
+	rng := collections.NewPrefixUntilTripleRange[time.Time, sdk.AccAddress, sdk.AccAddress](exp)
 	count := 0
 
 	keysToRemove := []collections.Triple[time.Time, sdk.AccAddress, sdk.AccAddress]{}
diff --git a/x/staking/keeper/delegation.go b/x/staking/keeper/delegation.go
index 18211132265f..69ad3843ee9c 100644
--- a/x/staking/keeper/delegation.go
+++ b/x/staking/keeper/delegation.go
@@ -421,7 +421,7 @@ func (k Keeper) GetRedelegations(ctx context.Context, delegator sdk.AccAddress,
 	redelegations = make([]types.Redelegation, maxRetrieve)
 
 	i := 0
-	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator, false)
+	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator)
 	err = k.Redelegations.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], redelegation types.Redelegation) (stop bool, err error) {
 		if i >= int(maxRetrieve) {
 			return true, nil
@@ -441,7 +441,7 @@ func (k Keeper) GetRedelegations(ctx context.Context, delegator sdk.AccAddress,
 // GetRedelegationsFromSrcValidator returns all redelegations from a particular
 // validator.
 func (k Keeper) GetRedelegationsFromSrcValidator(ctx context.Context, valAddr sdk.ValAddress) (reds []types.Redelegation, err error) {
-	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](valAddr, false)
+	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](valAddr)
 	err = k.RedelegationsByValSrc.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], value []byte) (stop bool, err error) {
 		valSrcAddr, delAddr, valDstAddr := key.K1(), key.K2(), key.K3()
 
@@ -462,7 +462,7 @@ func (k Keeper) GetRedelegationsFromSrcValidator(ctx context.Context, valAddr sd
 
 // HasReceivingRedelegation checks if validator is receiving a redelegation.
 func (k Keeper) HasReceivingRedelegation(ctx context.Context, delAddr sdk.AccAddress, valDstAddr sdk.ValAddress) (bool, error) {
-	rng := collections.NewSuperPrefixedTripleRange[[]byte, []byte, []byte](valDstAddr, delAddr, false)
+	rng := collections.NewSuperPrefixedTripleRange[[]byte, []byte, []byte](valDstAddr, delAddr)
 	hasReceivingRedelegation := false
 	err := k.RedelegationsByValDst.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], value []byte) (stop bool, err error) {
 		hasReceivingRedelegation = true
diff --git a/x/staking/keeper/query_utils.go b/x/staking/keeper/query_utils.go
index ebe23b81ebed..625984dbd44a 100644
--- a/x/staking/keeper/query_utils.go
+++ b/x/staking/keeper/query_utils.go
@@ -107,7 +107,7 @@ func (k Keeper) GetAllRedelegations(
 	dstValFilter := !(dstValAddress.Empty())
 
 	redelegations := []types.Redelegation{}
-	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator, false)
+	rng := collections.NewPrefixedTripleRange[[]byte, []byte, []byte](delegator)
 	err := k.Redelegations.Walk(ctx, rng,
 		func(key collections.Triple[[]byte, []byte, []byte], redelegation types.Redelegation) (stop bool, err error) {
 			valSrcAddr, valDstAddr := key.K2(), key.K3()

From bc0f0607bb11b8e1588430bd38041f183061626b Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Mon, 2 Sep 2024 17:29:27 +0300
Subject: [PATCH 5/8] changing reverse iterator option to be of Option type and
 not bool

---
 collections/triple.go      | 71 ++++++++++++++++++++++----------------
 collections/triple_test.go |  4 +--
 2 files changed, 43 insertions(+), 32 deletions(-)

diff --git a/collections/triple.go b/collections/triple.go
index cb4d8aac2a8e..70595c5bd1e4 100644
--- a/collections/triple.go
+++ b/collections/triple.go
@@ -273,58 +273,69 @@ func (t tripleKeyCodec[K1, K2, K3]) SizeNonTerminal(key Triple[K1, K2, K3]) int
 	return size
 }
 
+// defaultConfig has all the options disabled, except Color and TimeFormat
+var defaultConfig = Config{
+	Order: OrderAscending,
+}
+
+// Config defines configuration for the Triple collection iterator.
+type Config struct {
+	Order Order
+}
+
+type Option func(*Config)
+
+// OrderOption sets the order for the iterator.
+func OrderOption(order Order) Option {
+	return func(cfg *Config) {
+		cfg.Order = order
+	}
+}
+
 // NewPrefixUntilTripleRange defines a collection query which ranges until the provided Pair prefix.
-// Unstable: this API might change in the future. Use the first boolean option to get a reverse iterator (reverse=true).
-func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1, opts ...bool) Ranger[Triple[K1, K2, K3]] {
-	key := TriplePrefix[K1, K2, K3](k1)
-	order := OrderAscending
-	if len(opts) > 0 {
-		reverse := opts[0]
-		if reverse {
-			order = OrderDescending
-		}
+// Unstable: this API might change in the future.
+func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1, options ...Option) Ranger[Triple[K1, K2, K3]] {
+	cfg := defaultConfig
+	for _, opt := range options {
+		opt(&cfg)
 	}
 
+	key := TriplePrefix[K1, K2, K3](k1)
+
 	return &Range[Triple[K1, K2, K3]]{
 		end:   RangeKeyPrefixEnd(key),
-		order: order,
+		order: cfg.Order,
 	}
 }
 
-// NewPrefixedTripleRange provides a Range for all keys prefixed with the given
-// first part of the Triple key. Use the first boolean option to get a reverse iterator (reverse=true).
-func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1, opts ...bool) Ranger[Triple[K1, K2, K3]] {
-	key := TriplePrefix[K1, K2, K3](k1)
-	order := OrderAscending
-	if len(opts) > 0 {
-		reverse := opts[0]
-		if reverse {
-			order = OrderDescending
-		}
+// NewPrefixedTripleRange provides a Range for all keys prefixed with the given first part of the Triple key.
+func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1, options ...Option) Ranger[Triple[K1, K2, K3]] {
+	cfg := defaultConfig
+	for _, opt := range options {
+		opt(&cfg)
 	}
 
+	key := TriplePrefix[K1, K2, K3](k1)
+
 	return &Range[Triple[K1, K2, K3]]{
 		start: RangeKeyExact(key),
 		end:   RangeKeyPrefixEnd(key),
-		order: order,
+		order: cfg.Order,
 	}
 }
 
 // NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given
-// first and second parts of the Triple key. Use the first boolean option to get a reverse iterator (reverse=true).
-func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2, opts ...bool) Ranger[Triple[K1, K2, K3]] {
+// first and second parts of the Triple key.
+func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2, options ...Option) Ranger[Triple[K1, K2, K3]] {
 	key := TripleSuperPrefix[K1, K2, K3](k1, k2)
-	order := OrderAscending
-	if len(opts) > 0 {
-		reverse := opts[0]
-		if reverse {
-			order = OrderDescending
-		}
+	cfg := defaultConfig
+	for _, opt := range options {
+		opt(&cfg)
 	}
 
 	return &Range[Triple[K1, K2, K3]]{
 		start: RangeKeyExact(key),
 		end:   RangeKeyPrefixEnd(key),
-		order: order,
+		order: cfg.Order,
 	}
 }
diff --git a/collections/triple_test.go b/collections/triple_test.go
index a6dae56abfee..d4ae4778f53c 100644
--- a/collections/triple_test.go
+++ b/collections/triple_test.go
@@ -46,7 +46,7 @@ func TestTripleRange(t *testing.T) {
 	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.NewPrefixedTripleRange[uint64, string, []byte](uint64(1), true))
+	iter, err = keySet.Iterate(ctx, collections.NewPrefixedTripleRange[uint64, string, []byte](uint64(1), collections.OrderOption(collections.OrderDescending)))
 	require.NoError(t, err)
 	gotKeys, err = iter.Keys()
 	require.NoError(t, err)
@@ -63,7 +63,7 @@ func TestTripleRange(t *testing.T) {
 	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.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A", true))
+	iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A", collections.OrderOption(collections.OrderDescending)))
 	require.NoError(t, err)
 	gotKeys, err = iter.Keys()
 	require.NoError(t, err)

From 7e970f2834a5957c6138f6ed80e75d8cb4ed97d1 Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Tue, 3 Sep 2024 14:53:53 +0300
Subject: [PATCH 6/8] reverted order option and created separate functions for
 reverse iterator to avoid API breakage

---
 collections/triple.go      | 71 ++++++++++++++++++--------------------
 collections/triple_test.go |  4 +--
 2 files changed, 35 insertions(+), 40 deletions(-)

diff --git a/collections/triple.go b/collections/triple.go
index 70595c5bd1e4..6ab0dee49829 100644
--- a/collections/triple.go
+++ b/collections/triple.go
@@ -273,69 +273,64 @@ func (t tripleKeyCodec[K1, K2, K3]) SizeNonTerminal(key Triple[K1, K2, K3]) int
 	return size
 }
 
-// defaultConfig has all the options disabled, except Color and TimeFormat
-var defaultConfig = Config{
-	Order: OrderAscending,
+// NewPrefixUntilTripleRange defines a collection query which ranges until the provided Pair prefix.
+// Unstable: this API might change in the future.
+func NewPrefixUntilTripleRange[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),
+	}
 }
 
-// Config defines configuration for the Triple collection iterator.
-type Config struct {
-	Order Order
+// NewPrefixedTripleRange provides a Range for all keys prefixed with the given
+// first part of the Triple key.
+func NewPrefixedTripleRange[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),
+	}
 }
 
-type Option func(*Config)
-
-// OrderOption sets the order for the iterator.
-func OrderOption(order Order) Option {
-	return func(cfg *Config) {
-		cfg.Order = order
+// NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given
+// first and second parts of the Triple key.
+func NewSuperPrefixedTripleRange[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),
 	}
 }
 
-// NewPrefixUntilTripleRange defines a collection query which ranges until the provided Pair prefix.
+// NewPrefixUntilTripleRangeReversed defines a collection query which ranges until the provided Pair prefix
+// in reverse order.
 // Unstable: this API might change in the future.
-func NewPrefixUntilTripleRange[K1, K2, K3 any](k1 K1, options ...Option) Ranger[Triple[K1, K2, K3]] {
-	cfg := defaultConfig
-	for _, opt := range options {
-		opt(&cfg)
-	}
-
+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: cfg.Order,
+		order: OrderDescending,
 	}
 }
 
-// NewPrefixedTripleRange provides a Range for all keys prefixed with the given first part of the Triple key.
-func NewPrefixedTripleRange[K1, K2, K3 any](k1 K1, options ...Option) Ranger[Triple[K1, K2, K3]] {
-	cfg := defaultConfig
-	for _, opt := range options {
-		opt(&cfg)
-	}
-
+// 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: cfg.Order,
+		order: OrderDescending,
 	}
 }
 
 // NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given
-// first and second parts of the Triple key.
-func NewSuperPrefixedTripleRange[K1, K2, K3 any](k1 K1, k2 K2, options ...Option) Ranger[Triple[K1, K2, K3]] {
+// 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)
-	cfg := defaultConfig
-	for _, opt := range options {
-		opt(&cfg)
-	}
-
 	return &Range[Triple[K1, K2, K3]]{
 		start: RangeKeyExact(key),
 		end:   RangeKeyPrefixEnd(key),
-		order: cfg.Order,
+		order: OrderDescending,
 	}
 }
diff --git a/collections/triple_test.go b/collections/triple_test.go
index d4ae4778f53c..c605a7aeaba7 100644
--- a/collections/triple_test.go
+++ b/collections/triple_test.go
@@ -46,7 +46,7 @@ func TestTripleRange(t *testing.T) {
 	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.NewPrefixedTripleRange[uint64, string, []byte](uint64(1), collections.OrderOption(collections.OrderDescending)))
+	iter, err = keySet.Iterate(ctx, collections.NewPrefixedTripleRangeReversed[uint64, string, []byte](uint64(1)))
 	require.NoError(t, err)
 	gotKeys, err = iter.Keys()
 	require.NoError(t, err)
@@ -63,7 +63,7 @@ func TestTripleRange(t *testing.T) {
 	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.NewSuperPrefixedTripleRange[uint64, string, []byte](1, "A", collections.OrderOption(collections.OrderDescending)))
+	iter, err = keySet.Iterate(ctx, collections.NewSuperPrefixedTripleRangeReversed[uint64, string, []byte](1, "A"))
 	require.NoError(t, err)
 	gotKeys, err = iter.Keys()
 	require.NoError(t, err)

From a56e94317477e96b81982f69e82989ea91adb699 Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Sat, 12 Oct 2024 15:56:54 +0300
Subject: [PATCH 7/8] used unsafe bytes conversion to compare ref keys, added
 warning comment to RefKeys()

---
 collections/indexes/multi.go | 25 +++++++++++++------------
 collections/triple.go        |  4 ++--
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/collections/indexes/multi.go b/collections/indexes/multi.go
index 6a4a7f2d1a22..484e7485a8f5 100644
--- a/collections/indexes/multi.go
+++ b/collections/indexes/multi.go
@@ -2,9 +2,9 @@ package indexes
 
 import (
 	"context"
+	"crypto/sha256"
 	"errors"
-	"fmt"
-	"reflect"
+	"unsafe"
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/collections/codec"
@@ -125,31 +125,32 @@ func (m *Multi[ReferenceKey, PrimaryKey, Value]) MatchExact(ctx context.Context,
 
 // 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) {
-	// 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)
+	iter, err := m.refKeys.IterateRaw(ctx, nil, nil, collections.OrderAscending)
 	if err != nil {
 		return nil, err
 	}
 
 	keys := []ReferenceKey{}
-	visited := map[interface{}]struct{}{}
+	visited := map[[32]byte]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 {
+			// compare the byte representation of ref keys
+			refKey := key.K1()
+			unsafeRefKey := *(*[]byte)(unsafe.Pointer(&refKey))
+
+			// use SHA256 hash as map keys
+			if _, ok := visited[sha256.Sum256(unsafeRefKey)]; ok {
 				continue
 			}
-			visited[refKey] = struct{}{}
+			visited[sha256.Sum256(unsafeRefKey)] = struct{}{}
 		}
 		keys = append(keys, key.K1())
 	}
diff --git a/collections/triple.go b/collections/triple.go
index d783ab045218..f83412059d8c 100644
--- a/collections/triple.go
+++ b/collections/triple.go
@@ -351,7 +351,7 @@ func NewPrefixUntilTripleRangeReversed[K1, K2, K3 any](k1 K1) Ranger[Triple[K1,
 	}
 }
 
-// NewPrefixedTripleRange provides a Range for all keys prefixed with the given
+// 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)
@@ -362,7 +362,7 @@ func NewPrefixedTripleRangeReversed[K1, K2, K3 any](k1 K1) Ranger[Triple[K1, K2,
 	}
 }
 
-// NewSuperPrefixedTripleRange provides a Range for all keys prefixed with the given
+// 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)

From 551ed3475e3b532544b0a2b054712882cedaf824 Mon Sep 17 00:00:00 2001
From: Oren <oren@lavanet.xyz>
Date: Mon, 21 Oct 2024 19:09:43 +0300
Subject: [PATCH 8/8] fix ref key encoding for unique keys map

---
 collections/indexes/multi.go | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/collections/indexes/multi.go b/collections/indexes/multi.go
index 484e7485a8f5..144753b782f0 100644
--- a/collections/indexes/multi.go
+++ b/collections/indexes/multi.go
@@ -2,9 +2,7 @@ package indexes
 
 import (
 	"context"
-	"crypto/sha256"
 	"errors"
-	"unsafe"
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/collections/codec"
@@ -134,25 +132,32 @@ func (m *Multi[ReferenceKey, PrimaryKey, Value]) RefKeys(ctx context.Context, un
 	}
 
 	keys := []ReferenceKey{}
-	visited := map[[32]byte]struct{}{}
+	visited := map[string]struct{}{}
 	for ; iter.Valid(); iter.Next() {
 		key, err := iter.Key()
 		if err != nil {
 			return nil, err
 		}
+		refKey := key.K1()
 
 		if unique {
-			// compare the byte representation of ref keys
-			refKey := key.K1()
-			unsafeRefKey := *(*[]byte)(unsafe.Pointer(&refKey))
+			// 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
+			}
 
-			// use SHA256 hash as map keys
-			if _, ok := visited[sha256.Sum256(unsafeRefKey)]; ok {
+			// check if visited
+			if _, ok := visited[string(buf)]; ok {
 				continue
 			}
-			visited[sha256.Sum256(unsafeRefKey)] = struct{}{}
+			visited[string(buf)] = struct{}{}
 		}
-		keys = append(keys, key.K1())
+		keys = append(keys, refKey)
 	}
 
 	return keys, nil