refactor(x/staking): migrate RedelegationByValDstIndexKey key to collections

x/staking/keeper/delegation.go

@@ -522,14 +522,17 @@ 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) {
-	store := k.storeService.OpenKVStore(ctx)
-	prefix := types.GetREDsByDelToValDstIndexKey(delAddr, valDstAddr)
-	iterator, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix))
+	rng := collections.NewSuperPrefixedTripleRange[[]byte, []byte, []byte](valDstAddr, delAddr)
+	hasAtleastOneEntry := false
+	err := k.RedelegationsByValDst.Walk(ctx, rng, func(key collections.Triple[[]byte, []byte, []byte], value []byte) (stop bool, err error) {
+		hasAtleastOneEntry = true
+		return true, nil // returning true here to stop the iterations after 1st finding
+	})
 	if err != nil {
 		return false, err
 	}
-	defer iterator.Close()
-	return iterator.Valid(), nil
+
+	return hasAtleastOneEntry, nil
 }
 
 // HasMaxRedelegationEntries checks if the redelegation entries reached maximum limit.
@@ -557,7 +560,6 @@ func (k Keeper) SetRedelegation(ctx context.Context, red types.Redelegation) err
 		return err
 	}
 
-	store := k.storeService.OpenKVStore(ctx)
 	valSrcAddr, err := k.validatorAddressCodec.StringToBytes(red.ValidatorSrcAddress)
 	if err != nil {
 		return err
@@ -575,7 +577,7 @@ func (k Keeper) SetRedelegation(ctx context.Context, red types.Redelegation) err
 		return err
 	}
 
-	return store.Set(types.GetREDByValDstIndexKey(delegatorAddress, valSrcAddr, valDestAddr), []byte{})
+	return k.RedelegationsByValDst.Set(ctx, collections.Join3(valDestAddr, delegatorAddress, valSrcAddr), []byte{})
 }
 
 // SetRedelegationEntry adds an entry to the unbonding delegation at the given
@@ -645,7 +647,6 @@ func (k Keeper) RemoveRedelegation(ctx context.Context, red types.Redelegation)
 		return err
 	}
 
-	store := k.storeService.OpenKVStore(ctx)
 	valSrcAddr, err := k.validatorAddressCodec.StringToBytes(red.ValidatorSrcAddress)
 	if err != nil {
 		return err
@@ -663,7 +664,7 @@ func (k Keeper) RemoveRedelegation(ctx context.Context, red types.Redelegation)
 		return err
 	}
 
-	return store.Delete(types.GetREDByValDstIndexKey(delegatorAddress, valSrcAddr, valDestAddr))
+	return k.RedelegationsByValDst.Remove(ctx, collections.Join3(valDestAddr, delegatorAddress, valSrcAddr))
 }
 
 // redelegation queue timeslice operations

x/staking/keeper/keeper.go

@@ -44,6 +44,7 @@ type Keeper struct {
 	Redelegations               collections.Map[collections.Triple[[]byte, []byte, []byte], types.Redelegation]
 	Delegations                 collections.Map[collections.Pair[sdk.AccAddress, sdk.ValAddress], types.Delegation]
 	UnbondingIndex              collections.Map[uint64, []byte]
+	RedelegationsByValDst       collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
 	RedelegationsByValSrc       collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
 }
 
@@ -110,6 +111,7 @@ func NewKeeper(
 			collcodec.KeyToValueCodec(sdk.ValAddressKey),
 		),
 		UnbondingType: collections.NewMap(sb, types.UnbondingTypeKey, "unbonding_type", collections.Uint64Key, collections.Uint64Value),
+		// key format is: 52 | lengthPrefixedBytes(AccAddr) | lengthPrefixedBytes(SrcValAddr) | lengthPrefixedBytes(DstValAddr)
 		Redelegations: collections.NewMap(
 			sb, types.RedelegationKey,
 			"redelegations",
@@ -121,7 +123,7 @@ func NewKeeper(
 			codec.CollValue[types.Redelegation](cdc),
 		),
 		UnbondingIndex: collections.NewMap(sb, types.UnbondingIndexKey, "unbonding_index", collections.Uint64Key, collections.BytesValue),
-		// key format is: 53 | lengthPrefixedBytes(DstValAddr) | lengthPrefixedBytes(AccAddr) | lengthPrefixedBytes(SrcValAddr)
+		// key format is: 53 | lengthPrefixedBytes(SrcValAddr) | lengthPrefixedBytes(AccAddr) | lengthPrefixedBytes(DstValAddr)
 		RedelegationsByValSrc: collections.NewMap(
 			sb, types.RedelegationByValSrcIndexKey,
 			"redelegations_by_val_src",
@@ -132,6 +134,17 @@ func NewKeeper(
 			),
 			collections.BytesValue,
 		),
+		// key format is: 54 | lengthPrefixedBytes(DstValAddr) | lengthPrefixedBytes(AccAddr) | lengthPrefixedBytes(SrcValAddr)
+		RedelegationsByValDst: collections.NewMap(
+			sb, types.RedelegationByValDstIndexKey,
+			"redelegations_by_val_dst",
+			collections.TripleKeyCodec(
+				collections.BytesKey,
+				collections.BytesKey,
+				sdk.LengthPrefixedBytesKey, // sdk.LengthPrefixedBytesKey is needed to retain state compatibility
+			),
+			collections.BytesValue,
+		),
 	}
 
 	schema, err := sb.Build()

x/staking/keeper/keeper_test.go

@@ -113,6 +113,23 @@ func (s *KeeperTestSuite) TestLastTotalPower() {
 	require.True(expTotalPower.Equal(resTotalPower))
 }
 
+// getREDByValDstIndexKey creates the index-key for a redelegation, stored by destination-validator-index
+// VALUE: none (key rearrangement used)
+func getREDByValDstIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte {
+	REDSToValsDstKey := getREDsToValDstIndexKey(valDstAddr)
+	offset := len(REDSToValsDstKey)
+
+	// key is of the form REDSToValsDstKey || delAddrLen (1 byte) || delAddr || valSrcAddrLen (1 byte) || valSrcAddr
+	key := make([]byte, offset+2+len(delAddr)+len(valSrcAddr))
+	copy(key[0:offset], REDSToValsDstKey)
+	key[offset] = byte(len(delAddr))
+	copy(key[offset+1:offset+1+len(delAddr)], delAddr.Bytes())
+	key[offset+1+len(delAddr)] = byte(len(valSrcAddr))
+	copy(key[offset+2+len(delAddr):], valSrcAddr.Bytes())
+
+	return key
+}
+
 // GetREDByValSrcIndexKey creates the index-key for a redelegation, stored by source-validator-index
 // VALUE: none (key rearrangement used)
 func getREDByValSrcIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte {
@@ -130,14 +147,21 @@ func getREDByValSrcIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.V
 	return key
 }
 
+// GetREDsToValDstIndexKey returns a key prefix for indexing a redelegation to a
+// destination (target) validator.
+func getREDsToValDstIndexKey(valDstAddr sdk.ValAddress) []byte {
+	redelegationByValDstIndexKey := []byte{0x36}
+	return append(redelegationByValDstIndexKey, addresstypes.MustLengthPrefix(valDstAddr)...)
+}
+
 // GetREDsFromValSrcIndexKey returns a key prefix for indexing a redelegation to
 // a source validator.
 func getREDsFromValSrcIndexKey(valSrcAddr sdk.ValAddress) []byte {
 	redelegationByValSrcIndexKey := []byte{0x35}
 	return append(redelegationByValSrcIndexKey, addresstypes.MustLengthPrefix(valSrcAddr)...)
 }
 
-func (s *KeeperTestSuite) TestRedelegationsMigrationToColls() {
+func (s *KeeperTestSuite) TestSrcRedelegationsMigrationToColls() {
 	s.SetupTest()
 
 	addrs, valAddrs := createValAddrs(101)
@@ -169,6 +193,38 @@ func (s *KeeperTestSuite) TestRedelegationsMigrationToColls() {
 	s.Require().NoError(err)
 }
 
+func (s *KeeperTestSuite) TestDstRedelegationsMigrationToColls() {
+	s.SetupTest()
+
+	addrs, valAddrs := createValAddrs(101)
+
+	err := testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			// legacy method to set in the state
+			s.ctx.KVStore(s.key).Set(getREDByValDstIndexKey(addrs[i], valAddrs[i], valAddrs[i+1]), []byte{})
+		},
+		"4beb77994beff3c8ad9cecca9ee3a74fb551356250f0b8bd3936c4e4f506443b", // this hash obtained when ran this test in main branch
+	)
+	s.Require().NoError(err)
+
+	err = testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			// using collections
+			err := s.stakingKeeper.RedelegationsByValDst.Set(s.ctx, collections.Join3(valAddrs[i+1].Bytes(), addrs[i].Bytes(), valAddrs[i].Bytes()), []byte{})
+			s.Require().NoError(err)
+		},
+		"4beb77994beff3c8ad9cecca9ee3a74fb551356250f0b8bd3936c4e4f506443b",
+	)
+
+	s.Require().NoError(err)
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/staking/migrations/v2/keys.go

@@ -13,11 +13,14 @@ const (
 )
 
 var (
-	ValidatorsByConsAddrKey      = []byte{0x22} // prefix for validators by consensus address
+	ValidatorsByConsAddrKey = []byte{0x22} // prefix for validators by consensus address
+
 	DelegationKey                = []byte{0x31} // prefix for the delegation
 	RedelegationKey              = []byte{0x34} // key for a redelegation
+	RedelegationByValDstIndexKey = []byte{0x36} // prefix for each key for an redelegation, by destination validator operator
 	RedelegationByValSrcIndexKey = []byte{0x35} // prefix for each key for an redelegation, by source validator operator
-	HistoricalInfoKey            = []byte{0x50} // prefix for the historical info
+
+	HistoricalInfoKey = []byte{0x50} // prefix for the historical info
 )
 
 // GetHistoricalInfoKey returns a key prefix for indexing HistoricalInfo objects.
@@ -63,6 +66,23 @@ func GetREDsKey(delAddr sdk.AccAddress) []byte {
 	return append(RedelegationKey, address.MustLengthPrefix(delAddr)...)
 }
 
+// GetREDByValDstIndexKey creates the index-key for a redelegation, stored by destination-validator-index
+// VALUE: none (key rearrangement used)
+func GetREDByValDstIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte {
+	REDSToValsDstKey := GetREDsToValDstIndexKey(valDstAddr)
+	offset := len(REDSToValsDstKey)
+
+	// key is of the form REDSToValsDstKey || delAddrLen (1 byte) || delAddr || valSrcAddrLen (1 byte) || valSrcAddr
+	key := make([]byte, offset+2+len(delAddr)+len(valSrcAddr))
+	copy(key[0:offset], REDSToValsDstKey)
+	key[offset] = byte(len(delAddr))
+	copy(key[offset+1:offset+1+len(delAddr)], delAddr.Bytes())
+	key[offset+1+len(delAddr)] = byte(len(valSrcAddr))
+	copy(key[offset+2+len(delAddr):], valSrcAddr.Bytes())
+
+	return key
+}
+
 // GetREDByValSrcIndexKey creates the index-key for a redelegation, stored by source-validator-index
 // VALUE: none (key rearrangement used)
 func GetREDByValSrcIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte {
@@ -80,6 +100,12 @@ func GetREDByValSrcIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.V
 	return key
 }
 
+// GetREDsToValDstIndexKey returns a key prefix for indexing a redelegation to a
+// destination (target) validator.
+func GetREDsToValDstIndexKey(valDstAddr sdk.ValAddress) []byte {
+	return append(RedelegationByValDstIndexKey, address.MustLengthPrefix(valDstAddr)...)
+}
+
 // GetREDsFromValSrcIndexKey returns a key prefix for indexing a redelegation to
 // a source validator.
 func GetREDsFromValSrcIndexKey(valSrcAddr sdk.ValAddress) []byte {

x/staking/migrations/v2/store_test.go

@@ -95,7 +95,7 @@ func TestStoreMigration(t *testing.T) {
 		{
 			"RedelegationByValDstIndexKey",
 			v1.GetREDByValDstIndexKey(addr4, valAddr1, valAddr2),
-			types.GetREDByValDstIndexKey(addr4, valAddr1, valAddr2),
+			v2.GetREDByValDstIndexKey(addr4, valAddr1, valAddr2),
 		},
 		{
 			"UnbondingQueueKey",

x/staking/types/keys.go

@@ -44,7 +44,7 @@ var (
 	UnbondingDelegationByValIndexKey = []byte{0x33}              // prefix for each key for an unbonding-delegation, by validator operator
 	RedelegationKey                  = collections.NewPrefix(52) // key for a redelegation
 	RedelegationByValSrcIndexKey     = collections.NewPrefix(53) // prefix for each key for an redelegation, by source validator operator
-	RedelegationByValDstIndexKey     = []byte{0x36}              // prefix for each key for an redelegation, by destination validator operator
+	RedelegationByValDstIndexKey     = collections.NewPrefix(54) // prefix for each key for an redelegation, by destination validator operator
 
 	UnbondingIDKey    = collections.NewPrefix(55) // key for the counter for the incrementing id for UnbondingOperations
 	UnbondingIndexKey = collections.NewPrefix(56) // prefix for an index for looking up unbonding operations by their IDs
@@ -247,41 +247,6 @@ func GetREDKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []
 	return key
 }
 
-// GetREDByValDstIndexKey creates the index-key for a redelegation, stored by destination-validator-index
-// VALUE: none (key rearrangement used)
-func GetREDByValDstIndexKey(delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) []byte {
-	REDSToValsDstKey := GetREDsToValDstIndexKey(valDstAddr)
-	offset := len(REDSToValsDstKey)
-
-	// key is of the form REDSToValsDstKey || delAddrLen (1 byte) || delAddr || valSrcAddrLen (1 byte) || valSrcAddr
-	key := make([]byte, offset+2+len(delAddr)+len(valSrcAddr))
-	copy(key[0:offset], REDSToValsDstKey)
-	key[offset] = byte(len(delAddr))
-	copy(key[offset+1:offset+1+len(delAddr)], delAddr.Bytes())
-	key[offset+1+len(delAddr)] = byte(len(valSrcAddr))
-	copy(key[offset+2+len(delAddr):], valSrcAddr.Bytes())
-
-	return key
-}
-
-// GetREDKeyFromValDstIndexKey rearranges the ValDstIndexKey to get the REDKey
-func GetREDKeyFromValDstIndexKey(indexKey []byte) []byte {
-	// note that first byte is prefix byte, which we remove
-	kv.AssertKeyAtLeastLength(indexKey, 2)
-	addrs := indexKey[1:]
-
-	valDstAddrLen := addrs[0]
-	kv.AssertKeyAtLeastLength(addrs, int(valDstAddrLen)+2)
-	valDstAddr := addrs[1 : valDstAddrLen+1]
-	delAddrLen := addrs[valDstAddrLen+1]
-	kv.AssertKeyAtLeastLength(addrs, int(valDstAddrLen)+int(delAddrLen)+3)
-	delAddr := addrs[valDstAddrLen+2 : valDstAddrLen+2+delAddrLen]
-	kv.AssertKeyAtLeastLength(addrs, int(valDstAddrLen)+int(delAddrLen)+4)
-	valSrcAddr := addrs[valDstAddrLen+delAddrLen+3:]
-
-	return GetREDKey(delAddr, valSrcAddr, valDstAddr)
-}
-
 // GetRedelegationTimeKey returns a key prefix for indexing an unbonding
 // redelegation based on a completion time.
 func GetRedelegationTimeKey(timestamp time.Time) []byte {
@@ -294,15 +259,3 @@ func GetRedelegationTimeKey(timestamp time.Time) []byte {
 func GetREDsKey(delAddr sdk.AccAddress) []byte {
 	return append(RedelegationKey, address.MustLengthPrefix(delAddr)...)
 }
-
-// GetREDsToValDstIndexKey returns a key prefix for indexing a redelegation to a
-// destination (target) validator.
-func GetREDsToValDstIndexKey(valDstAddr sdk.ValAddress) []byte {
-	return append(RedelegationByValDstIndexKey, address.MustLengthPrefix(valDstAddr)...)
-}
-
-// GetREDsByDelToValDstIndexKey returns a key prefix for indexing a redelegation
-// from an address to a source validator.
-func GetREDsByDelToValDstIndexKey(delAddr sdk.AccAddress, valDstAddr sdk.ValAddress) []byte {
-	return append(GetREDsToValDstIndexKey(valDstAddr), address.MustLengthPrefix(delAddr)...)
-}

x/staking/types/keys_test.go

@@ -19,11 +19,7 @@ import (
 
 var (
 	keysPK1   = ed25519.GenPrivKeyFromSecret([]byte{1}).PubKey()
-	keysPK2   = ed25519.GenPrivKeyFromSecret([]byte{2}).PubKey()
-	keysPK3   = ed25519.GenPrivKeyFromSecret([]byte{3}).PubKey()
 	keysAddr1 = keysPK1.Address()
-	keysAddr2 = keysPK2.Address()
-	keysAddr3 = keysPK3.Address()
 )
 
 func TestGetValidatorPowerRank(t *testing.T) {
@@ -52,33 +48,6 @@ func TestGetValidatorPowerRank(t *testing.T) {
 	}
 }
 
-func TestGetREDByValDstIndexKey(t *testing.T) {
-	tests := []struct {
-		delAddr    sdk.AccAddress
-		valSrcAddr sdk.ValAddress
-		valDstAddr sdk.ValAddress
-		wantHex    string
-	}{
-		{
-			sdk.AccAddress(keysAddr1), sdk.ValAddress(keysAddr1), sdk.ValAddress(keysAddr1),
-			"361463d771218209d8bd03c482f69dfba57310f086091463d771218209d8bd03c482f69dfba57310f086091463d771218209d8bd03c482f69dfba57310f08609",
-		},
-		{
-			sdk.AccAddress(keysAddr1), sdk.ValAddress(keysAddr2), sdk.ValAddress(keysAddr3),
-			"36143ab62f0d93849be495e21e3e9013a517038f45bd1463d771218209d8bd03c482f69dfba57310f08609145ef3b5f25c54946d4a89fc0d09d2f126614540f2",
-		},
-		{
-			sdk.AccAddress(keysAddr2), sdk.ValAddress(keysAddr1), sdk.ValAddress(keysAddr3),
-			"36143ab62f0d93849be495e21e3e9013a517038f45bd145ef3b5f25c54946d4a89fc0d09d2f126614540f21463d771218209d8bd03c482f69dfba57310f08609",
-		},
-	}
-	for i, tt := range tests {
-		got := hex.EncodeToString(types.GetREDByValDstIndexKey(tt.delAddr, tt.valSrcAddr, tt.valDstAddr))
-
-		require.Equal(t, tt.wantHex, got, "Keys did not match on test case %d", i)
-	}
-}
-
 func TestGetValidatorQueueKey(t *testing.T) {
 	ts := time.Now()
 	height := int64(1024)