refactor(x/staking): migrate delegations by validator to use collecti…

…ons (#17155)

CHANGELOG.md

@@ -39,7 +39,9 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ## [Unreleased]
 
 ### Features
+
 * (x/bank) [#16795](https://github.com/cosmos/cosmos-sdk/pull/16852) Add `DenomMetadataByQueryString` query in bank module to support metadata query by query string.
+
 ### Improvements
 
 * (x/group, x/gov) [#17109](https://github.com/cosmos/cosmos-sdk/pull/17109) Let proposal summary be 40x longer than metadata limit.
@@ -53,7 +55,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
-* (x/slashing) [17063](https://github.com/cosmos/cosmos-sdk/pull/17063) Use collections for `HistoricalInfo`:
+* (x/staking) [17063](https://github.com/cosmos/cosmos-sdk/pull/17063) Use collections for `HistoricalInfo`:
     * remove `Keeper`: `GetHistoricalInfo`, `SetHistoricalInfo`,
 * (x/staking) [17062](https://github.com/cosmos/cosmos-sdk/pull/17062) Use collections for `ValidatorUpdates`:
     * remove `Keeper`: `SetValidatorUpdates`, `GetValidatorUpdates`

tests/integration/staking/keeper/validator_bench_test.go

@@ -2,9 +2,11 @@ package keeper_test
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"testing"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
 	storetypes "cosmossdk.io/store/types"
 
@@ -155,29 +157,27 @@ func updateValidatorDelegations(f *fixture, existingValAddr, newValAddr sdk.ValA
 
 	store := f.sdkCtx.KVStore(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)
-		}
+	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)
 
 		// remove old operator addr from delegation
 		if err := k.RemoveDelegation(f.sdkCtx, delegation); err != nil {
-			panic(err)
+			return true, err
 		}
 
 		delegation.ValidatorAddress = newValAddr.String()
 		// add with new operator addr
 		if err := k.SetDelegation(f.sdkCtx, delegation); err != nil {
-			panic(err)
+			return true, err
 		}
 
+		return false, nil
+	})
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		panic(err)
 	}
 }

x/staking/keeper/delegation.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"time"
 
+	"cosmossdk.io/collections"
 	corestore "cosmossdk.io/core/store"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/math"
@@ -65,32 +66,29 @@ func (k Keeper) GetAllDelegations(ctx context.Context) (delegations []types.Dele
 
 // GetValidatorDelegations returns all delegations to a specific validator.
 // Useful for querier.
-func (k Keeper) GetValidatorDelegations(ctx context.Context, valAddr sdk.ValAddress) (delegations []types.Delegation, err error) {
+func (k Keeper) GetValidatorDelegations(ctx context.Context, valAddr sdk.ValAddress) ([]types.Delegation, error) {
 	store := k.storeService.OpenKVStore(ctx)
-	prefix := types.GetDelegationsByValPrefixKey(valAddr)
-	iterator, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix))
-	if err != nil {
-		return delegations, err
-	}
-	defer iterator.Close()
 
-	for ; iterator.Valid(); iterator.Next() {
+	var delegations []types.Delegation
+	rng := collections.NewPrefixedPairRange[sdk.ValAddress, sdk.AccAddress](valAddr)
+	err := k.DelegationsByValidator.Walk(ctx, rng, func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], _ []byte) (stop bool, err error) {
 		var delegation types.Delegation
-		valAddr, delAddr, err := types.ParseDelegationsByValKey(iterator.Key())
-		if err != nil {
-			return delegations, err
-		}
-
+		valAddr, delAddr := key.K1(), key.K2()
 		bz, err := store.Get(types.GetDelegationKey(delAddr, valAddr))
 		if err != nil {
-			return delegations, err
+			return true, err
 		}
 
 		if err := k.cdc.Unmarshal(bz, &delegation); err != nil {
-			return delegations, err
+			return true, err
 		}
 
 		delegations = append(delegations, delegation)
+
+		return false, nil
+	})
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		return nil, err
 	}
 
 	return delegations, nil
@@ -142,7 +140,7 @@ func (k Keeper) SetDelegation(ctx context.Context, delegation types.Delegation)
 	}
 
 	// set the delegation in validator delegator index
-	return store.Set(types.GetDelegationsByValKey(valAddr, delegatorAddress), []byte{})
+	return k.DelegationsByValidator.Set(ctx, collections.Join(sdk.ValAddress(valAddr), sdk.AccAddress(delegatorAddress)), []byte{})
 }
 
 // RemoveDelegation removes a delegation
@@ -168,7 +166,7 @@ func (k Keeper) RemoveDelegation(ctx context.Context, delegation types.Delegatio
 		return err
 	}
 
-	return store.Delete(types.GetDelegationsByValKey(valAddr, delegatorAddress))
+	return k.DelegationsByValidator.Remove(ctx, collections.Join(sdk.ValAddress(valAddr), sdk.AccAddress(delegatorAddress)))
 }
 
 // GetUnbondingDelegations returns a given amount of all the delegator unbonding-delegations.

x/staking/keeper/grpc_query.go

@@ -7,6 +7,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/store/prefix"
 	storetypes "cosmossdk.io/store/types"
 
@@ -101,24 +102,26 @@ func (k Querier) ValidatorDelegations(ctx context.Context, req *types.QueryValid
 	}
 
 	store := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
-	delStore := prefix.NewStore(store, types.GetDelegationsByValPrefixKey(valAddr))
 
 	var (
 		dels    types.Delegations
 		pageRes *query.PageResponse
 	)
-	pageRes, err = query.Paginate(delStore, req.Pagination, func(delAddr, value []byte) error {
-		bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
 
-		var delegation types.Delegation
-		err = k.cdc.Unmarshal(bz, &delegation)
-		if err != nil {
-			return err
-		}
+	dels, pageRes, err = query.CollectionPaginate(ctx, k.DelegationsByValidator,
+		req.Pagination, func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], _ []byte) (types.Delegation, error) {
+			valAddr, delAddr := key.K1(), key.K2()
+			var delegation types.Delegation
+
+			bz := store.Get(types.GetDelegationKey(delAddr, valAddr))
+			err = k.cdc.Unmarshal(bz, &delegation)
+			if err != nil {
+				return types.Delegation{}, err
+			}
+
+			return delegation, nil
+		}, query.WithCollectionPaginationPairPrefix[sdk.ValAddress, sdk.AccAddress](valAddr))
 
-		dels = append(dels, delegation)
-		return nil
-	})
 	if err != nil {
 		delegations, pageResponse, err := k.getValidatorDelegationsLegacy(ctx, req)
 		if err != nil {

x/staking/keeper/keeper.go

@@ -32,10 +32,11 @@ type Keeper struct {
 	validatorAddressCodec addresscodec.Codec
 	consensusAddressCodec addresscodec.Codec
 
-	Schema           collections.Schema
-	HistoricalInfo   collections.Map[uint64, types.HistoricalInfo]
-	LastTotalPower   collections.Item[math.Int]
-	ValidatorUpdates collections.Item[types.ValidatorUpdates]
+	Schema                 collections.Schema
+	HistoricalInfo         collections.Map[uint64, types.HistoricalInfo]
+	LastTotalPower         collections.Item[math.Int]
+	ValidatorUpdates       collections.Item[types.ValidatorUpdates]
+	DelegationsByValidator collections.Map[collections.Pair[sdk.ValAddress, sdk.AccAddress], []byte]
 }
 
 // NewKeeper creates a new staking Keeper instance
@@ -79,6 +80,12 @@ func NewKeeper(
 		LastTotalPower:        collections.NewItem(sb, types.LastTotalPowerKey, "last_total_power", sdk.IntValue),
 		HistoricalInfo:        collections.NewMap(sb, types.HistoricalInfoKey, "historical_info", collections.Uint64Key, codec.CollValue[types.HistoricalInfo](cdc)),
 		ValidatorUpdates:      collections.NewItem(sb, types.ValidatorUpdatesKey, "validator_updates", codec.CollValue[types.ValidatorUpdates](cdc)),
+		DelegationsByValidator: collections.NewMap(
+			sb, types.DelegationByValIndexKey,
+			"delegations_by_validator",
+			collections.PairKeyCodec(sdk.LengthPrefixedAddressKey(sdk.ValAddressKey), sdk.AccAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+			collections.BytesValue,
+		),
 	}
 
 	schema, err := sb.Build()

x/staking/migrations/v5/keys.go

@@ -71,3 +71,33 @@ func GetDelegationsByValPrefixKey(valAddr sdk.ValAddress) []byte {
 func GetDelegationsByValKey(valAddr sdk.ValAddress, delAddr sdk.AccAddress) []byte {
 	return append(GetDelegationsByValPrefixKey(valAddr), delAddr...)
 }
+
+// ParseDelegationsByValKey parses given key and returns validator, delegator address bytes
+func ParseDelegationsByValKey(bz []byte) (sdk.ValAddress, sdk.AccAddress, error) {
+	prefixLength := len(DelegationByValIndexKey)
+	if prefix := bz[:prefixLength]; !bytes.Equal(prefix, DelegationByValIndexKey) {
+		return nil, nil, fmt.Errorf("invalid prefix; expected: %X, got: %x", DelegationByValIndexKey, prefix)
+	}
+
+	bz = bz[prefixLength:] // remove the prefix byte
+	if len(bz) == 0 {
+		return nil, nil, fmt.Errorf("no bytes left to parse: %X", bz)
+	}
+
+	valAddrLen := bz[0]
+	bz = bz[1:] // remove the length byte of validator address.
+	if len(bz) == 0 {
+		return nil, nil, fmt.Errorf("no bytes left to parse validator address: %X", bz)
+	}
+
+	val := bz[0:int(valAddrLen)]
+
+	bz = bz[int(valAddrLen):] // remove the delegator bytes
+	if len(bz) == 0 {
+		return nil, nil, fmt.Errorf("no bytes left to parse delegator address: %X", bz)
+	}
+
+	del := bz
+
+	return val, del, nil
+}

x/staking/migrations/v5/migrations_test.go

@@ -115,7 +115,7 @@ func getValDelegations(ctx sdk.Context, cdc codec.Codec, storeKey storetypes.Sto
 	iterator := storetypes.KVStorePrefixIterator(store, v5.GetDelegationsByValPrefixKey(valAddr))
 	for ; iterator.Valid(); iterator.Next() {
 		var delegation stakingtypes.Delegation
-		valAddr, delAddr, err := stakingtypes.ParseDelegationsByValKey(iterator.Key())
+		valAddr, delAddr, err := v5.ParseDelegationsByValKey(iterator.Key())
 		if err != nil {
 			panic(err)
 		}

x/staking/types/keys.go

@@ -215,47 +215,6 @@ func GetDelegationKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
 	return append(GetDelegationsKey(delAddr), address.MustLengthPrefix(valAddr)...)
 }
 
-// GetDelegationsByValKey creates the key for delegations by validator address
-// VALUE: staking/Delegation
-func GetDelegationsByValKey(valAddr sdk.ValAddress, delAddr sdk.AccAddress) []byte {
-	return append(GetDelegationsByValPrefixKey(valAddr), delAddr...)
-}
-
-// GetDelegationsByValPrefixKey builds a prefix key bytes with the given validator address bytes.
-func GetDelegationsByValPrefixKey(valAddr sdk.ValAddress) []byte {
-	return append(DelegationByValIndexKey, address.MustLengthPrefix(valAddr)...)
-}
-
-// ParseDelegationsByValKey parses given key and returns validator, delegator address bytes
-func ParseDelegationsByValKey(bz []byte) (sdk.ValAddress, sdk.AccAddress, error) {
-	prefixLength := len(DelegationByValIndexKey)
-	if prefix := bz[:prefixLength]; !bytes.Equal(prefix, DelegationByValIndexKey) {
-		return nil, nil, fmt.Errorf("invalid prefix; expected: %X, got: %x", DelegationByValIndexKey, prefix)
-	}
-
-	bz = bz[prefixLength:] // remove the prefix byte
-	if len(bz) == 0 {
-		return nil, nil, fmt.Errorf("no bytes left to parse: %X", bz)
-	}
-
-	valAddrLen := bz[0]
-	bz = bz[1:] // remove the length byte of validator address.
-	if len(bz) == 0 {
-		return nil, nil, fmt.Errorf("no bytes left to parse validator address: %X", bz)
-	}
-
-	val := bz[0:int(valAddrLen)]
-
-	bz = bz[int(valAddrLen):] // remove the delegator bytes
-	if len(bz) == 0 {
-		return nil, nil, fmt.Errorf("no bytes left to parse delegator address: %X", bz)
-	}
-
-	del := bz
-
-	return val, del, nil
-}
-
 // GetDelegationsKey creates the prefix for a delegator for all validators
 func GetDelegationsKey(delAddr sdk.AccAddress) []byte {
 	return append(DelegationKey, address.MustLengthPrefix(delAddr)...)