refactor(x/staking): migrate redelegation key to use collections

CHANGELOG.md

@@ -70,10 +70,12 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (x/staking) [#17315](https://github.com/cosmos/cosmos-sdk/pull/17044) Use collections for `RedelegationKey`:
+    * remove from `keeper`: `GetRedelegation`
 * (types) `module.BeginBlockAppModule` has been replaced by Core API `appmodule.HasBeginBlocker`.
 * (types) `module.EndBlockAppModule` has been replaced by Core API `appmodule.HasEndBlocker` or `module.HasABCIEndBlock` when needing validator updates.
 * (types) [#17358](https://github.com/cosmos/cosmos-sdk/pull/17358) Remove deprecated `sdk.Handler`, use `baseapp.MsgServiceHandler` instead.
-* (x/slashing) [17044](https://github.com/cosmos/cosmos-sdk/pull/17044) Use collections for `AddrPubkeyRelation`:
+* (x/slashing) [#17044](https://github.com/cosmos/cosmos-sdk/pull/17044) Use collections for `AddrPubkeyRelation`:
     * remove from `types`: `AddrPubkeyRelationKey`
     * remove from `Keeper`: `AddPubkey`
 * (x/staking) [#17260](https://github.com/cosmos/cosmos-sdk/pull/17260) Use collections for `DelegationKey`:

tests/integration/staking/keeper/grpc_query_test.go

@@ -820,7 +820,7 @@ func TestGRPCQueryRedelegations(t *testing.T) {
 	assert.NilError(t, err)
 	applyValidatorSetUpdates(t, ctx, f.stakingKeeper, -1)
 
-	redel, found := f.stakingKeeper.GetRedelegation(ctx, addrAcc1, val1bz, val2bz)
+	redel, found := f.stakingKeeper.Redelegations.Get(ctx, collections.Join3(addrAcc1, valAddrs[0], valAddrs[1]))
 	assert.Assert(t, found)
 
 	var req *types.QueryRedelegationsRequest

tests/integration/staking/keeper/slash_test.go

@@ -157,7 +157,7 @@ func TestSlashRedelegation(t *testing.T) {
 	slashAmount, err = f.stakingKeeper.SlashRedelegation(f.sdkCtx, validator, rd, 0, fraction)
 	assert.NilError(t, err)
 	assert.Assert(t, slashAmount.Equal(math.NewInt(5)))
-	rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
+	rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	assert.Assert(t, found)
 	assert.Assert(t, len(rd.Entries) == 1)
 
@@ -438,7 +438,7 @@ func TestSlashWithRedelegation(t *testing.T) {
 	oldBonded = f.bankKeeper.GetBalance(f.sdkCtx, bondedPool.GetAddress(), bondDenom).Amount
 
 	// read updating redelegation
-	rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
+	rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	assert.Assert(t, found)
 	assert.Assert(t, len(rd.Entries) == 1)
 	// read updated validator
@@ -475,7 +475,7 @@ func TestSlashWithRedelegation(t *testing.T) {
 	oldBonded = f.bankKeeper.GetBalance(f.sdkCtx, bondedPool.GetAddress(), bondDenom).Amount
 
 	// read updating redelegation
-	rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
+	rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	assert.Assert(t, found)
 	assert.Assert(t, len(rd.Entries) == 1)
 	// read updated validator
@@ -506,7 +506,7 @@ func TestSlashWithRedelegation(t *testing.T) {
 	oldBonded = f.bankKeeper.GetBalance(f.sdkCtx, bondedPool.GetAddress(), bondDenom).Amount
 
 	// read updating redelegation
-	rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
+	rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	assert.Assert(t, found)
 	assert.Assert(t, len(rd.Entries) == 1)
 	// apply TM updates
@@ -536,7 +536,7 @@ func TestSlashWithRedelegation(t *testing.T) {
 	assert.Assert(math.IntEq(t, oldNotBonded, notBondedPoolBalance))
 
 	// read updating redelegation
-	rd, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
+	rd, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	assert.Assert(t, found)
 	assert.Assert(t, len(rd.Entries) == 1)
 	// read updated validator
@@ -607,7 +607,7 @@ func TestSlashBoth(t *testing.T) {
 	assert.Assert(math.IntEq(t, oldNotBonded.Sub(burnedNotBondedAmount), notBondedPoolBalance))
 
 	// read updating redelegation
-	rdA, found = f.stakingKeeper.GetRedelegation(f.sdkCtx, addrDels[0], addrVals[0], addrVals[1])
+	rdA, found = f.stakingKeeper.Redelegations.Get(f.sdkCtx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	assert.Assert(t, found)
 	assert.Assert(t, len(rdA.Entries) == 1)
 	// read updated validator

x/staking/keeper/delegation.go

@@ -293,22 +293,18 @@ func (k Keeper) IterateDelegatorDelegations(ctx context.Context, delegator sdk.A
 
 // IterateDelegatorRedelegations iterates through one delegator's redelegations.
 func (k Keeper) IterateDelegatorRedelegations(ctx context.Context, delegator sdk.AccAddress, cb func(red types.Redelegation) (stop bool)) error {
-	store := k.storeService.OpenKVStore(ctx)
-	delegatorPrefixKey := types.GetREDsKey(delegator)
-	iterator, err := store.Iterator(delegatorPrefixKey, storetypes.PrefixEndBytes(delegatorPrefixKey))
+	rng := collections.NewPrefixedTripleRange[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress](delegator)
+	err := k.Redelegations.Walk(ctx, rng, func(key collections.Triple[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress], red types.Redelegation) (stop bool, err error) {
+		if cb(red) {
+			return true, nil
+		}
+
+		return false, nil
+	})
 	if err != nil {
 		return err
 	}
 
-	for ; iterator.Valid(); iterator.Next() {
-		red, err := types.UnmarshalRED(k.cdc, iterator.Value())
-		if err != nil {
-			return err
-		}
-		if cb(red) {
-			break
-		}
-	}
 	return nil
 }
 
@@ -500,41 +496,23 @@ func (k Keeper) DequeueAllMatureUBDQueue(ctx context.Context, currTime time.Time
 func (k Keeper) GetRedelegations(ctx context.Context, delegator sdk.AccAddress, maxRetrieve uint16) (redelegations []types.Redelegation, err error) {
 	redelegations = make([]types.Redelegation, maxRetrieve)
 
-	store := k.storeService.OpenKVStore(ctx)
-	delegatorPrefixKey := types.GetREDsKey(delegator)
-	iterator, err := store.Iterator(delegatorPrefixKey, storetypes.PrefixEndBytes(delegatorPrefixKey))
-	if err != nil {
-		return nil, err
-	}
-
 	i := 0
-	for ; iterator.Valid() && i < int(maxRetrieve); iterator.Next() {
-		redelegation, err := types.UnmarshalRED(k.cdc, iterator.Value())
-		if err != nil {
-			return nil, err
+	rng := collections.NewPrefixedTripleRange[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress](delegator)
+	err = k.Redelegations.Walk(ctx, rng, func(key collections.Triple[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress], redelegation types.Redelegation) (stop bool, err error) {
+		if i >= int(maxRetrieve) {
+			return true, nil
 		}
+
 		redelegations[i] = redelegation
 		i++
-	}
-
-	return redelegations[:i], nil // trim if the array length < maxRetrieve
-}
-
-// GetRedelegation returns a redelegation.
-func (k Keeper) GetRedelegation(ctx context.Context, delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress) (red types.Redelegation, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-	key := types.GetREDKey(delAddr, valSrcAddr, valDstAddr)
+		return false, nil
+	})
 
-	value, err := store.Get(key)
 	if err != nil {
-		return red, err
-	}
-
-	if value == nil {
-		return red, types.ErrNoRedelegation
+		return nil, err
 	}
 
-	return types.UnmarshalRED(k.cdc, value)
+	return redelegations[:i], nil // trim if the array length < maxRetrieve
 }
 
 // GetRedelegationsFromSrcValidator returns all redelegations from a particular
@@ -578,9 +556,9 @@ func (k Keeper) HasReceivingRedelegation(ctx context.Context, delAddr sdk.AccAdd
 
 // HasMaxRedelegationEntries checks if the redelegation entries reached maximum limit.
 func (k Keeper) HasMaxRedelegationEntries(ctx context.Context, delegatorAddr sdk.AccAddress, validatorSrcAddr, validatorDstAddr sdk.ValAddress) (bool, error) {
-	red, err := k.GetRedelegation(ctx, delegatorAddr, validatorSrcAddr, validatorDstAddr)
+	red, err := k.Redelegations.Get(ctx, collections.Join3(delegatorAddr, validatorSrcAddr, validatorDstAddr))
 	if err != nil {
-		if err == types.ErrNoRedelegation {
+		if errors.Is(err, collections.ErrNotFound) {
 			return false, nil
 		}
 
@@ -602,7 +580,6 @@ func (k Keeper) SetRedelegation(ctx context.Context, red types.Redelegation) err
 	}
 
 	store := k.storeService.OpenKVStore(ctx)
-	bz := types.MustMarshalRED(k.cdc, red)
 	valSrcAddr, err := k.validatorAddressCodec.StringToBytes(red.ValidatorSrcAddress)
 	if err != nil {
 		return err
@@ -611,8 +588,8 @@ func (k Keeper) SetRedelegation(ctx context.Context, red types.Redelegation) err
 	if err != nil {
 		return err
 	}
-	key := types.GetREDKey(delegatorAddress, valSrcAddr, valDestAddr)
-	if err = store.Set(key, bz); err != nil {
+
+	if err = k.Redelegations.Set(ctx, collections.Join3(sdk.AccAddress(delegatorAddress), sdk.ValAddress(valSrcAddr), sdk.ValAddress(valDestAddr)), red); err != nil {
 		return err
 	}
 
@@ -636,10 +613,10 @@ func (k Keeper) SetRedelegationEntry(ctx context.Context,
 		return types.Redelegation{}, err
 	}
 
-	red, err := k.GetRedelegation(ctx, delegatorAddr, validatorSrcAddr, validatorDstAddr)
+	red, err := k.Redelegations.Get(ctx, collections.Join3(delegatorAddr, validatorSrcAddr, validatorDstAddr))
 	if err == nil {
 		red.AddEntry(creationHeight, minTime, balance, sharesDst, id)
-	} else if errors.Is(err, types.ErrNoRedelegation) {
+	} else if errors.Is(err, collections.ErrNotFound) {
 		red = types.NewRedelegation(delegatorAddr, validatorSrcAddr,
 			validatorDstAddr, creationHeight, minTime, balance, sharesDst, id, k.validatorAddressCodec, k.authKeeper.AddressCodec())
 	} else {
@@ -665,22 +642,19 @@ func (k Keeper) SetRedelegationEntry(ctx context.Context,
 
 // IterateRedelegations iterates through all redelegations.
 func (k Keeper) IterateRedelegations(ctx context.Context, fn func(index int64, red types.Redelegation) (stop bool)) error {
-	store := k.storeService.OpenKVStore(ctx)
-	iterator, err := store.Iterator(types.RedelegationKey, storetypes.PrefixEndBytes(types.RedelegationKey))
-	if err != nil {
-		return err
-	}
-	defer iterator.Close()
+	var i int64
+	err := k.Redelegations.Walk(ctx, nil,
+		func(key collections.Triple[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress], red types.Redelegation) (bool, error) {
+			if stop := fn(i, red); stop {
+				return true, nil
+			}
+			i++
 
-	for i := int64(0); iterator.Valid(); iterator.Next() {
-		red, err := types.UnmarshalRED(k.cdc, iterator.Value())
-		if err != nil {
-			return err
-		}
-		if stop := fn(i, red); stop {
-			break
-		}
-		i++
+			return false, nil
+		},
+	)
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		return err
 	}
 
 	return nil
@@ -702,8 +676,8 @@ func (k Keeper) RemoveRedelegation(ctx context.Context, red types.Redelegation)
 	if err != nil {
 		return err
 	}
-	redKey := types.GetREDKey(delegatorAddress, valSrcAddr, valDestAddr)
-	if err = store.Delete(redKey); err != nil {
+
+	if err = k.Redelegations.Remove(ctx, collections.Join3(sdk.AccAddress(delegatorAddress), sdk.ValAddress(valSrcAddr), sdk.ValAddress(valDestAddr))); err != nil {
 		return err
 	}
 
@@ -1250,7 +1224,7 @@ func (k Keeper) BeginRedelegation(
 func (k Keeper) CompleteRedelegation(
 	ctx context.Context, delAddr sdk.AccAddress, valSrcAddr, valDstAddr sdk.ValAddress,
 ) (sdk.Coins, error) {
-	red, err := k.GetRedelegation(ctx, delAddr, valSrcAddr, valDstAddr)
+	red, err := k.Redelegations.Get(ctx, collections.Join3(delAddr, valSrcAddr, valDstAddr))
 	if err != nil {
 		return nil, err
 	}

x/staking/keeper/delegation_test.go

@@ -682,7 +682,7 @@ func (s *KeeperTestSuite) TestGetRedelegationsFromSrcValidator() {
 	// set and retrieve a record
 	err := keeper.SetRedelegation(ctx, rd)
 	require.NoError(err)
-	resBond, err := keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
+	resBond, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	require.NoError(err)
 
 	// get the redelegations one time
@@ -717,7 +717,7 @@ func (s *KeeperTestSuite) TestRedelegation() {
 	// set and retrieve a record
 	err = keeper.SetRedelegation(ctx, rd)
 	require.NoError(err)
-	resRed, err := keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
+	resRed, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	require.NoError(err)
 
 	redelegations, err := keeper.GetRedelegationsFromSrcValidator(ctx, addrVals[0])
@@ -735,7 +735,7 @@ func (s *KeeperTestSuite) TestRedelegation() {
 	require.Equal(1, len(redelegations))
 	require.Equal(redelegations[0], resRed)
 
-	// check if has the redelegation
+	// check if it has the redelegation
 	has, err = keeper.HasReceivingRedelegation(ctx, addrDels[0], addrVals[1])
 	require.NoError(err)
 	require.True(has)
@@ -745,7 +745,7 @@ func (s *KeeperTestSuite) TestRedelegation() {
 	err = keeper.SetRedelegation(ctx, rd)
 	require.NoError(err)
 
-	resRed, err = keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
+	resRed, err = keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
 	require.NoError(err)
 	require.Equal(rd, resRed)
 
@@ -762,8 +762,8 @@ func (s *KeeperTestSuite) TestRedelegation() {
 	// delete a record
 	err = keeper.RemoveRedelegation(ctx, rd)
 	require.NoError(err)
-	_, err = keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
-	require.ErrorIs(err, stakingtypes.ErrNoRedelegation)
+	_, err = keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
+	require.ErrorIs(err, collections.ErrNotFound)
 
 	redelegations, err = keeper.GetRedelegations(ctx, addrDels[0], 5)
 	require.NoError(err)
@@ -976,7 +976,7 @@ func (s *KeeperTestSuite) TestRedelegateFromUnbondingValidator() {
 	require.NoError(err)
 
 	// retrieve the unbonding delegation
-	ubd, err := keeper.GetRedelegation(ctx, addrDels[1], addrVals[0], addrVals[1])
+	ubd, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[1], addrVals[0], addrVals[1]))
 	require.NoError(err)
 	require.Len(ubd.Entries, 1)
 	require.Equal(blockHeight, ubd.Entries[0].CreationHeight)
@@ -1050,8 +1050,8 @@ func (s *KeeperTestSuite) TestRedelegateFromUnbondedValidator() {
 	require.NoError(err)
 
 	// no red should have been found
-	red, err := keeper.GetRedelegation(ctx, addrDels[0], addrVals[0], addrVals[1])
-	require.ErrorIs(err, stakingtypes.ErrNoRedelegation, "%v", red)
+	red, err := keeper.Redelegations.Get(ctx, collections.Join3(addrDels[0], addrVals[0], addrVals[1]))
+	require.ErrorIs(err, collections.ErrNotFound, "%v", red)
 }
 
 func (s *KeeperTestSuite) TestUnbondingDelegationAddEntry() {

x/staking/keeper/grpc_query.go

@@ -412,7 +412,7 @@ func (k Querier) Redelegations(ctx context.Context, req *types.QueryRedelegation
 	case req.DelegatorAddr == "" && req.SrcValidatorAddr != "" && req.DstValidatorAddr == "":
 		redels, pageRes, err = queryRedelegationsFromSrcValidator(store, k, req)
 	default:
-		redels, pageRes, err = queryAllRedelegations(store, k, req)
+		redels, pageRes, err = queryAllRedelegations(ctx, store, k, req)
 	}
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
@@ -505,7 +505,7 @@ func queryRedelegation(ctx context.Context, k Querier, req *types.QueryRedelegat
 		return nil, err
 	}
 
-	redel, err := k.GetRedelegation(ctx, delAddr, srcValAddr, dstValAddr)
+	redel, err := k.Keeper.Redelegations.Get(ctx, collections.Join3(sdk.AccAddress(delAddr), sdk.ValAddress(srcValAddr), sdk.ValAddress(dstValAddr)))
 	if err != nil {
 		return nil, status.Errorf(
 			codes.NotFound,
@@ -539,21 +539,18 @@ func queryRedelegationsFromSrcValidator(store storetypes.KVStore, k Querier, req
 	return redels, res, err
 }
 
-func queryAllRedelegations(store storetypes.KVStore, k Querier, req *types.QueryRedelegationsRequest) (redels types.Redelegations, res *query.PageResponse, err error) {
+func queryAllRedelegations(ctx context.Context, store storetypes.KVStore, k Querier, req *types.QueryRedelegationsRequest) (redels types.Redelegations, res *query.PageResponse, err error) {
 	delAddr, err := k.authKeeper.AddressCodec().StringToBytes(req.DelegatorAddr)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	redStore := prefix.NewStore(store, types.GetREDsKey(delAddr))
-	res, err = query.Paginate(redStore, req.Pagination, func(key, value []byte) error {
-		redelegation, err := types.UnmarshalRED(k.cdc, value)
-		if err != nil {
-			return err
-		}
-		redels = append(redels, redelegation)
-		return nil
-	})
+	redels, res, err = query.CollectionPaginate(ctx, k.Keeper.Redelegations, req.Pagination, func(_ collections.Triple[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress], red types.Redelegation) (types.Redelegation, error) {
+		return red, nil
+	}, query.WithCollectionPaginationTriplePrefix[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress](delAddr))
+	if err != nil {
+		return nil, nil, err
+	}
 
 	return redels, res, err
 }

x/staking/keeper/keeper.go

@@ -41,6 +41,7 @@ type Keeper struct {
 	UnbondingID                 collections.Sequence
 	ValidatorByConsensusAddress collections.Map[sdk.ConsAddress, sdk.ValAddress]
 	UnbondingType               collections.Map[uint64, uint64]
+	Redelegations               collections.Map[collections.Triple[sdk.AccAddress, sdk.ValAddress, sdk.ValAddress], types.Redelegation]
 	Delegations                 collections.Map[collections.Pair[sdk.AccAddress, sdk.ValAddress], types.Delegation]
 	UnbondingIndex              collections.Map[uint64, []byte]
 }
@@ -107,7 +108,17 @@ func NewKeeper(
 			sdk.LengthPrefixedAddressKey(sdk.ConsAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
 			collcodec.KeyToValueCodec(sdk.ValAddressKey),
 		),
-		UnbondingType:  collections.NewMap(sb, types.UnbondingTypeKey, "unbonding_type", collections.Uint64Key, collections.Uint64Value),
+		UnbondingType: collections.NewMap(sb, types.UnbondingTypeKey, "unbonding_type", collections.Uint64Key, collections.Uint64Value),
+		Redelegations: collections.NewMap(
+			sb, types.RedelegationKey,
+			"redelegations",
+			collections.TripleKeyCodec(
+				sdk.LengthPrefixedAddressKey(sdk.AccAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+				sdk.LengthPrefixedAddressKey(sdk.ValAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+				sdk.LengthPrefixedAddressKey(sdk.ValAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+			),
+			codec.CollValue[types.Redelegation](cdc),
+		),
 		UnbondingIndex: collections.NewMap(sb, types.UnbondingIndexKey, "unbonding_index", collections.Uint64Key, collections.BytesValue),
 	}