refactor: add delegations by validator index

tests/integration/staking/keeper/determinstic_test.go

@@ -322,7 +322,7 @@ func TestGRPCValidatorDelegations(t *testing.T) {
 		ValidatorAddr: validator.OperatorAddress,
 	}
 
-	testdata.DeterministicIterations(f.ctx, t, req, f.queryClient.ValidatorDelegations, 11985, false)
+	testdata.DeterministicIterations(f.ctx, t, req, f.queryClient.ValidatorDelegations, 14475, false)
 }
 
 func TestGRPCValidatorUnbondingDelegations(t *testing.T) {

tests/integration/staking/keeper/validator_bench_test.go

@@ -1,6 +1,13 @@
 package keeper_test
 
-import "testing"
+import (
+	"fmt"
+	"testing"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktestutil "github.com/cosmos/cosmos-sdk/x/bank/testutil"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+)
 
 func BenchmarkGetValidator(b *testing.B) {
 	// 900 is the max number we are allowed to use in order to avoid simtestutil.CreateTestPubKeys
@@ -27,3 +34,69 @@ func BenchmarkGetValidator(b *testing.B) {
 		}
 	}
 }
+
+func BenchmarkGetValidatorDelegations(b *testing.B) {
+	powersNumber := 10
+	var totalPower int64
+	powers := make([]int64, powersNumber)
+	for i := range powers {
+		powers[i] = int64(i)
+		totalPower += int64(i)
+	}
+
+	app, ctx, _, valAddrs, vals := initValidators(b, totalPower, len(powers), powers)
+	for _, validator := range vals {
+		app.StakingKeeper.SetValidator(ctx, validator)
+	}
+
+	delegationsNum := 1000
+	delegators := make([]sdk.AccAddress, delegationsNum)
+	for _, val := range valAddrs {
+		for i := 0; i < delegationsNum; i++ {
+			delegator := sdk.AccAddress(fmt.Sprintf("address%d", i))
+			banktestutil.FundAccount(app.BankKeeper, ctx, delegator,
+				sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(int64(i)))))
+			NewDel := types.NewDelegation(delegator, val, sdk.NewDec(int64(i)))
+			app.StakingKeeper.SetDelegation(ctx, NewDel)
+			delegators[i] = delegator
+		}
+	}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		app.StakingKeeper.UpdateValidatorDelegations(ctx, valAddrs[0], sdk.ValAddress("val"))
+	}
+}
+
+func BenchmarkGetValidatorDelegationsLegacy(b *testing.B) {
+	powersNumber := 10
+	var totalPower int64
+	powers := make([]int64, powersNumber)
+	for i := range powers {
+		powers[i] = int64(i)
+		totalPower += int64(i)
+	}
+
+	app, ctx, _, valAddrs, vals := initValidators(b, totalPower, len(powers), powers)
+	for _, validator := range vals {
+		app.StakingKeeper.SetValidator(ctx, validator)
+	}
+
+	delegationsNum := 1000
+	delegators := make([]sdk.AccAddress, delegationsNum)
+	for _, val := range valAddrs {
+		for i := 0; i < delegationsNum; i++ {
+			delegator := sdk.AccAddress(fmt.Sprintf("address%d", i))
+			banktestutil.FundAccount(app.BankKeeper, ctx, delegator,
+				sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(int64(i)))))
+			NewDel := types.NewDelegation(delegator, val, sdk.NewDec(int64(i)))
+			app.StakingKeeper.SetDelegation(ctx, NewDel)
+			delegators[i] = delegator
+		}
+	}
+
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		app.StakingKeeper.UpdateValidatorDelegationsLegacy(ctx, valAddrs[0], sdk.ValAddress("val"))
+	}
+}

x/staking/keeper/delegation.go

@@ -60,19 +60,70 @@ func (k Keeper) GetAllDelegations(ctx sdk.Context) (delegations []types.Delegati
 func (k Keeper) GetValidatorDelegations(ctx sdk.Context, valAddr sdk.ValAddress) (delegations []types.Delegation) {
 	store := ctx.KVStore(k.storeKey)
 
-	iterator := storetypes.KVStorePrefixIterator(store, types.DelegationKey)
+	iterator := storetypes.KVStorePrefixIterator(store, types.GetDelegationsByValPrefixKey(valAddr))
 	defer iterator.Close()
 
 	for ; iterator.Valid(); iterator.Next() {
-		delegation := types.MustUnmarshalDelegation(k.cdc, iterator.Value())
-		if delegation.GetValidatorAddr().Equals(valAddr) {
-			delegations = append(delegations, delegation)
+		var delegation types.Delegation
+		valAddr, delAddr, err := types.ParseDelegationsByValKey(iterator.Key())
+		if err != nil {
+			panic(err)
+		}
+
+		bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
+		if err := k.cdc.Unmarshal(bz, &delegation); err != nil {
+			panic(err)
 		}
+
+		delegations = append(delegations, delegation)
 	}
 
 	return delegations
 }
 
+func (k Keeper) UpdateValidatorDelegations(ctx sdk.Context, existingValAddr, newValAddr sdk.ValAddress) {
+	store := ctx.KVStore(k.storeKey)
+
+	itr := storetypes.KVStorePrefixIterator(store, types.GetDelegationsByValPrefixKey(existingValAddr))
+	defer itr.Close()
+
+	for ; itr.Valid(); itr.Next() {
+		key := itr.Key()
+		valAddr, delAddr, err := types.ParseDelegationsByValKey(key)
+		if err != nil {
+			panic(err)
+		}
+
+		bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
+		delegation := types.MustUnmarshalDelegation(k.cdc, bz)
+
+		// remove old operator addr from delegation
+		if err := k.RemoveDelegation(ctx, delegation); err != nil {
+			panic(err)
+		}
+
+		delegation.ValidatorAddress = newValAddr.String()
+		// add with new operator addr
+		k.SetDelegation(ctx, delegation)
+	}
+}
+
+func (k Keeper) UpdateValidatorDelegationsLegacy(ctx sdk.Context, existingValAddr, newValAddr sdk.ValAddress) {
+	store := ctx.KVStore(k.storeKey)
+
+	iterator := storetypes.KVStorePrefixIterator(store, types.DelegationKey)
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		delegation := types.MustUnmarshalDelegation(k.cdc, iterator.Value())
+		if delegation.GetValidatorAddr().Equals(existingValAddr) {
+			k.RemoveDelegation(ctx, delegation)
+			delegation.ValidatorAddress = newValAddr.String()
+			k.SetDelegation(ctx, delegation)
+		}
+	}
+}
+
 // GetDelegatorDelegations returns a given amount of all the delegations from a
 // delegator.
 func (k Keeper) GetDelegatorDelegations(ctx sdk.Context, delegator sdk.AccAddress, maxRetrieve uint16) (delegations []types.Delegation) {
@@ -103,6 +154,9 @@ func (k Keeper) SetDelegation(ctx sdk.Context, delegation types.Delegation) {
 	store := ctx.KVStore(k.storeKey)
 	b := types.MustMarshalDelegation(k.cdc, delegation)
 	store.Set(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()), b)
+
+	// set the delegation in validator delegator index
+	store.Set(types.GetDelegationsByValKey(delegation.GetValidatorAddr(), delegatorAddress), []byte{})
 }
 
 // RemoveDelegation removes a delegation
@@ -119,6 +173,8 @@ func (k Keeper) RemoveDelegation(ctx sdk.Context, delegation types.Delegation) e
 
 	store := ctx.KVStore(k.storeKey)
 	store.Delete(types.GetDelegationKey(delegatorAddress, delegation.GetValidatorAddr()))
+	store.Delete(types.GetDelegationsByValKey(delegation.GetValidatorAddr(), delegatorAddress))
+
 	return nil
 }
 

x/staking/keeper/delegation_test.go

@@ -144,6 +144,83 @@ func (s *KeeperTestSuite) TestDelegation() {
 	require.Equal(0, len(resBonds))
 }
 
+func (s *KeeperTestSuite) TestDelegationsByValIndex() {
+	ctx, keeper := s.ctx, s.stakingKeeper
+	require := s.Require()
+
+	addrDels, valAddrs := createValAddrs(3)
+
+	for _, addr := range addrDels {
+		s.accountKeeper.EXPECT().StringToBytes(addr.String()).Return(addr, nil).AnyTimes()
+		s.accountKeeper.EXPECT().BytesToString(addr).Return(addr.String(), nil).AnyTimes()
+		s.bankKeeper.EXPECT().DelegateCoinsFromAccountToModule(gomock.Any(), addr, gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+	}
+
+	// construct the validators
+	amts := []math.Int{sdk.NewInt(9), sdk.NewInt(8), sdk.NewInt(7)}
+	var validators [3]stakingtypes.Validator
+	for i, amt := range amts {
+		validators[i] = testutil.NewValidator(s.T(), valAddrs[i], PKs[i])
+		validators[i], _ = validators[i].AddTokensFromDel(amt)
+
+		validators[i] = stakingkeeper.TestingUpdateValidator(keeper, ctx, validators[i], true)
+	}
+
+	// delegate 2 tokens
+	//
+	// total delegations after delegating: del1 -> 2stake
+	_, err := s.msgServer.Delegate(ctx, stakingtypes.NewMsgDelegate(addrDels[0], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(2))))
+	require.NoError(err)
+
+	dels := s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 1)
+
+	// delegate 4 tokens
+	//
+	// total delegations after delegating: del1 -> 2stake, del2 -> 4stake
+	_, err = s.msgServer.Delegate(ctx, stakingtypes.NewMsgDelegate(addrDels[1], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(4))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 2)
+
+	// undelegate 1 token from del1
+	//
+	// total delegations after undelegating: del1 -> 1stake, del2 -> 4stake
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[0], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(1))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 2)
+
+	// undelegate 1 token from del1
+	//
+	// total delegations after undelegating: del2 -> 4stake
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[0], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(1))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 1)
+
+	// undelegate 2 tokens from del2
+	//
+	// total delegations after undelegating: del2 -> 2stake
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[1], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(2))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 1)
+
+	// undelegate 2 tokens from del2
+	//
+	// total delegations after undelegating: []
+	_, err = s.msgServer.Undelegate(ctx, stakingtypes.NewMsgUndelegate(addrDels[1], valAddrs[0], sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(2))))
+	require.NoError(err)
+
+	dels = s.stakingKeeper.GetValidatorDelegations(ctx, valAddrs[0])
+	require.Len(dels, 0)
+}
+
 // tests Get/Set/Remove UnbondingDelegation
 func (s *KeeperTestSuite) TestUnbondingDelegation() {
 	ctx, keeper := s.ctx, s.stakingKeeper

x/staking/keeper/grpc_query.go

@@ -92,33 +92,34 @@ func (k Querier) ValidatorDelegations(c context.Context, req *types.QueryValidat
 	if req.ValidatorAddr == "" {
 		return nil, status.Error(codes.InvalidArgument, "validator address cannot be empty")
 	}
+
+	valAddr, err := sdk.ValAddressFromBech32(req.ValidatorAddr)
+	if err != nil {
+		return nil, err
+	}
 	ctx := sdk.UnwrapSDKContext(c)
 
 	store := ctx.KVStore(k.storeKey)
-	valStore := prefix.NewStore(store, types.DelegationKey)
-	delegations, pageRes, err := query.GenericFilteredPaginate(k.cdc, valStore, req.Pagination, func(key []byte, delegation *types.Delegation) (*types.Delegation, error) {
-		valAddr, err := sdk.ValAddressFromBech32(req.ValidatorAddr)
-		if err != nil {
-			return nil, err
-		}
+	delStore := prefix.NewStore(store, types.GetDelegationsByValPrefixKey(valAddr))
 
-		if !delegation.GetValidatorAddr().Equals(valAddr) {
-			return nil, nil
+	var dels types.Delegations
+	pageRes, err := query.Paginate(delStore, req.Pagination, func(del, value []byte) error {
+		bz := store.Get(types.GetDelegationKey(del, valAddr))
+
+		var delegation types.Delegation
+		k.cdc.Unmarshal(bz, &delegation)
+		if err != nil {
+			return err
 		}
 
-		return delegation, nil
-	}, func() *types.Delegation {
-		return &types.Delegation{}
+		dels = append(dels, delegation)
+		return nil
 	})
+
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
 
-	dels := types.Delegations{}
-	for _, d := range delegations {
-		dels = append(dels, *d)
-	}
-
 	delResponses, err := DelegationsToDelegationResponses(ctx, k.Keeper, dels)
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())

x/staking/migrations/v5/keys.go

@@ -0,0 +1,11 @@
+package v5
+
+const (
+	// ModuleName is the name of the module
+	ModuleName = "staking"
+)
+
+var (
+	DelegationKey           = []byte{0x31} // key for a delegation
+	DelegationByValIndexKey = []byte{0x37} // key for delegations by a validator
+)

x/staking/migrations/v5/store.go

@@ -0,0 +1,27 @@
+package v5
+
+import (
+	storetypes "cosmossdk.io/store/types"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/staking/types"
+)
+
+func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.BinaryCodec) (storetypes.KVStore, error) {
+	store := ctx.KVStore(storeKey)
+
+	iterator := storetypes.KVStorePrefixIterator(store, DelegationKey)
+
+	for ; iterator.Valid(); iterator.Next() {
+		key := iterator.Key()
+		del, val, err := types.ParseDelegationKey(key)
+		if err != nil {
+			return store, err
+		}
+
+		store.Set(types.GetDelegationsByValKey(val, del), []byte{})
+	}
+
+	return store, nil
+}