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

tests/integration/distribution/keeper/msg_server_test.go

@@ -253,7 +253,7 @@ func TestMsgWithdrawDelegatorReward(t *testing.T) {
 				ValidatorAddress: f.valAddr.String(),
 			},
 			expErr:    true,
-			expErrMsg: "no delegation for (address, validator) tuple",
+			expErrMsg: "not found",
 		},
 		{
 			name: "validator with no delegations",

tests/integration/evidence/keeper/infraction_test.go

@@ -11,6 +11,7 @@ import (
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"gotest.tools/v3/assert"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/core/comet"
 	"cosmossdk.io/log"
@@ -243,7 +244,7 @@ func TestHandleDoubleSign(t *testing.T) {
 
 	// require we be able to unbond now
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
-	del, _ := f.stakingKeeper.GetDelegation(ctx, sdk.AccAddress(operatorAddr), operatorAddr)
+	del, _ := f.stakingKeeper.Delegations.Get(ctx, collections.Join(sdk.AccAddress(operatorAddr), operatorAddr))
 	validator, _ := f.stakingKeeper.GetValidator(ctx, operatorAddr)
 	totalBond := validator.TokensFromShares(del.GetShares()).TruncateInt()
 	tstaking.Ctx = ctx

tests/integration/staking/keeper/grpc_query_test.go

@@ -8,6 +8,7 @@ import (
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"gotest.tools/v3/assert"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
 
 	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
@@ -223,7 +224,7 @@ func TestGRPCQueryDelegatorValidator(t *testing.T) {
 				}
 			},
 			false,
-			"no delegation for (address, validator) tuple",
+			"not found",
 		},
 		{
 			"empty delegator address",
@@ -289,7 +290,7 @@ func TestGRPCQueryDelegation(t *testing.T) {
 	addrVal := vals[0].OperatorAddress
 	valAddr, err := sdk.ValAddressFromBech32(addrVal)
 	assert.NilError(t, err)
-	delegation, found := f.stakingKeeper.GetDelegation(ctx, addrAcc, valAddr)
+	delegation, found := f.stakingKeeper.Delegations.Get(ctx, collections.Join(addrAcc, valAddr))
 	assert.Assert(t, found)
 	var req *types.QueryDelegationRequest
 
@@ -358,7 +359,7 @@ func TestGRPCQueryDelegatorDelegations(t *testing.T) {
 	addrVal1 := vals[0].OperatorAddress
 	valAddr, err := sdk.ValAddressFromBech32(addrVal1)
 	assert.NilError(t, err)
-	delegation, found := f.stakingKeeper.GetDelegation(ctx, addrAcc, valAddr)
+	delegation, found := f.stakingKeeper.Delegations.Get(ctx, collections.Join(addrAcc, valAddr))
 	assert.Assert(t, found)
 	var req *types.QueryDelegatorDelegationsRequest
 
@@ -438,7 +439,7 @@ func TestGRPCQueryValidatorDelegations(t *testing.T) {
 	addrVal2 := valAddrs[4]
 	valAddr, err := sdk.ValAddressFromBech32(addrVal1)
 	assert.NilError(t, err)
-	delegation, found := f.stakingKeeper.GetDelegation(ctx, addrAcc, valAddr)
+	delegation, found := f.stakingKeeper.Delegations.Get(ctx, collections.Join(addrAcc, valAddr))
 	assert.Assert(t, found)
 
 	var req *types.QueryValidatorDelegationsRequest

tests/integration/staking/keeper/slash_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"gotest.tools/v3/assert"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
 
 	"github.com/cosmos/cosmos-sdk/codec/address"
@@ -167,7 +168,7 @@ func TestSlashRedelegation(t *testing.T) {
 	assert.DeepEqual(t, math.NewInt(10), rd.Entries[0].InitialBalance)
 
 	// shares decreased
-	del, found = f.stakingKeeper.GetDelegation(f.sdkCtx, addrDels[0], addrVals[1])
+	del, found = f.stakingKeeper.Delegations.Get(f.sdkCtx, collections.Join(addrDels[0], addrVals[1]))
 	assert.Assert(t, found)
 	assert.Equal(t, int64(5), del.Shares.RoundInt64())
 

tests/integration/staking/keeper/validator_bench_test.go

@@ -152,17 +152,16 @@ func updateValidatorDelegationsLegacy(f *fixture, existingValAddr, newValAddr sd
 }
 
 func updateValidatorDelegations(f *fixture, existingValAddr, newValAddr sdk.ValAddress) {
-	storeKey := f.keys[types.StoreKey]
-	cdc, k := f.cdc, f.stakingKeeper
-
-	store := f.sdkCtx.KVStore(storeKey)
+	k := f.stakingKeeper
 
 	rng := collections.NewPrefixedPairRange[sdk.ValAddress, sdk.AccAddress](existingValAddr)
 	err := k.DelegationsByValidator.Walk(f.sdkCtx, rng, func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], _ []byte) (stop bool, err error) {
 		valAddr, delAddr := key.K1(), key.K2()
 
-		bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
-		delegation := types.MustUnmarshalDelegation(cdc, bz)
+		delegation, err := k.Delegations.Get(f.sdkCtx, collections.Join(delAddr, valAddr))
+		if err != nil {
+			return true, err
+		}
 
 		// remove old operator addr from delegation
 		if err := k.RemoveDelegation(f.sdkCtx, delegation); err != nil {

x/staking/app_test.go

@@ -7,6 +7,7 @@ import (
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/stretchr/testify/require"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/depinject"
 	"cosmossdk.io/log"
 	"cosmossdk.io/math"
@@ -115,7 +116,7 @@ func TestStakingMsgs(t *testing.T) {
 
 	ctxCheck = app.BaseApp.NewContext(true)
 	require.True(t, sdk.Coins{genCoin.Sub(bondCoin)}.Equal(bankKeeper.GetAllBalances(ctxCheck, addr2)))
-	_, err = stakingKeeper.GetDelegation(ctxCheck, addr2, sdk.ValAddress(addr1))
+	_, err = stakingKeeper.Delegations.Get(ctxCheck, collections.Join(addr2, sdk.ValAddress(addr1)))
 	require.NoError(t, err)
 
 	// begin unbonding
@@ -126,8 +127,8 @@ func TestStakingMsgs(t *testing.T) {
 
 	// delegation should exist anymore
 	ctxCheck = app.BaseApp.NewContext(true)
-	_, err = stakingKeeper.GetDelegation(ctxCheck, addr2, sdk.ValAddress(addr1))
-	require.ErrorIs(t, err, types.ErrNoDelegation)
+	_, err = stakingKeeper.Delegations.Get(ctxCheck, collections.Join(addr2, sdk.ValAddress(addr1)))
+	require.ErrorIs(t, err, collections.ErrNotFound)
 
 	// balance should be the same because bonding not yet complete
 	require.True(t, sdk.Coins{genCoin.Sub(bondCoin)}.Equal(bankKeeper.GetAllBalances(ctxCheck, addr2)))

x/staking/keeper/alias_functions.go

@@ -2,7 +2,9 @@ package keeper
 
 import (
 	"context"
+	"errors"
 
+	"cosmossdk.io/collections"
 	storetypes "cosmossdk.io/store/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -115,7 +117,7 @@ func (k Keeper) GetValidatorSet() types.ValidatorSet {
 
 // Delegation gets the delegation interface for a particular set of delegator and validator addresses
 func (k Keeper) Delegation(ctx context.Context, addrDel sdk.AccAddress, addrVal sdk.ValAddress) (types.DelegationI, error) {
-	bond, err := k.GetDelegation(ctx, addrDel, addrVal)
+	bond, err := k.Delegations.Get(ctx, collections.Join(addrDel, addrVal))
 	if err != nil {
 		return nil, err
 	}
@@ -127,25 +129,19 @@ func (k Keeper) Delegation(ctx context.Context, addrDel sdk.AccAddress, addrVal
 func (k Keeper) IterateDelegations(ctx context.Context, delAddr sdk.AccAddress,
 	fn func(index int64, del types.DelegationI) (stop bool),
 ) error {
-	store := k.storeService.OpenKVStore(ctx)
-	delegatorPrefixKey := types.GetDelegationsKey(delAddr)
-	iterator, err := store.Iterator(delegatorPrefixKey, storetypes.PrefixEndBytes(delegatorPrefixKey))
-	if err != nil {
-		return err
-	}
-	defer iterator.Close()
-
-	for i := int64(0); iterator.Valid(); iterator.Next() {
-		del, err := types.UnmarshalDelegation(k.cdc, iterator.Value())
-		if err != nil {
-			return err
-		}
-
-		stop := fn(i, del)
+	var i int64
+	rng := collections.NewPrefixedPairRange[sdk.AccAddress, sdk.ValAddress](delAddr)
+	err := k.Delegations.Walk(ctx, rng, func(key collections.Pair[sdk.AccAddress, sdk.ValAddress], del types.Delegation) (stop bool, err error) {
+		stop = fn(i, del)
 		if stop {
-			break
+			return true, err
 		}
 		i++
+
+		return false, nil
+	})
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		return err
 	}
 
 	return nil