refactor(x/staking): Migrate UnbondingDelegationByValIndex to use collections

x/staking/keeper/delegation.go

@@ -215,26 +215,29 @@ func (k Keeper) GetUnbondingDelegation(ctx context.Context, delAddr sdk.AccAddre
 // particular validator.
 func (k Keeper) GetUnbondingDelegationsFromValidator(ctx context.Context, valAddr sdk.ValAddress) (ubds []types.UnbondingDelegation, err error) {
 	store := k.storeService.OpenKVStore(ctx)
-	prefix := types.GetUBDsByValIndexKey(valAddr)
-	iterator, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix))
-	if err != nil {
+	rng := collections.NewPrefixUntilPairRange[sdk.ValAddress, sdk.AccAddress](valAddr)
+	err = k.UnbondingDelegationByValIndex.Walk(
+		ctx,
+		rng,
+		func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], value []byte) (stop bool, err error) {
+			valAddr := key.K1()
+			delAddr := key.K2()
+			ubdkey := types.GetUBDKey(delAddr, valAddr)
+			ubdValue, err := store.Get(ubdkey)
+			if err != nil {
+				return true, err
+			}
+			unbondingDelegation, err := types.UnmarshalUBD(k.cdc, ubdValue)
+			if err != nil {
+				return true, err
+			}
+			ubds = append(ubds, unbondingDelegation)
+			return false, nil
+		},
+	)
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
 		return ubds, err
 	}
-	defer iterator.Close()
-
-	for ; iterator.Valid(); iterator.Next() {
-		key := types.GetUBDKeyFromValIndexKey(iterator.Key())
-		value, err := store.Get(key)
-		if err != nil {
-			return ubds, err
-		}
-		ubd, err := types.UnmarshalUBD(k.cdc, value)
-		if err != nil {
-			return ubds, err
-		}
-		ubds = append(ubds, ubd)
-	}
-
 	return ubds, nil
 }
 

x/staking/keeper/grpc_query.go

@@ -183,19 +183,25 @@ func (k Querier) ValidatorUnbondingDelegations(ctx context.Context, req *types.Q
 	}
 
 	store := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
-	srcValPrefix := types.GetUBDsByValIndexKey(valAddr)
-	ubdStore := prefix.NewStore(store, srcValPrefix)
-	pageRes, err := query.Paginate(ubdStore, req.Pagination, func(key, value []byte) error {
-		storeKey := types.GetUBDKeyFromValIndexKey(append(srcValPrefix, key...))
-		storeValue := store.Get(storeKey)
-
-		ubd, err := types.UnmarshalUBD(k.cdc, storeValue)
-		if err != nil {
-			return err
-		}
-		ubds = append(ubds, ubd)
-		return nil
-	})
+	_, pageRes, err := query.CollectionPaginate(
+		ctx,
+		k.UnbondingDelegationByValIndex,
+		req.Pagination,
+		func(key collections.Pair[sdk.ValAddress, sdk.AccAddress], value []byte) (types.UnbondingDelegation, error) {
+			valAddr := key.K1()
+			delAddr := key.K2()
+			ubdKey := types.GetUBDKey(delAddr, valAddr)
+			storeValue := store.Get(ubdKey)
+
+			ubd, err := types.UnmarshalUBD(k.cdc, storeValue)
+			if err != nil {
+				return types.UnbondingDelegation{}, err
+			}
+			ubds = append(ubds, ubd)
+			return ubd, nil
+		},
+		query.WithCollectionPaginationPairPrefix[sdk.ValAddress, sdk.AccAddress](valAddr),
+	)
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}

x/staking/keeper/keeper.go

@@ -33,14 +33,15 @@ 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]
-	DelegationsByValidator      collections.Map[collections.Pair[sdk.ValAddress, sdk.AccAddress], []byte]
-	UnbondingID                 collections.Sequence
-	ValidatorByConsensusAddress collections.Map[sdk.ConsAddress, sdk.ValAddress]
-	UnbondingType               collections.Map[uint64, uint64]
+	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]
+	UnbondingID                   collections.Sequence
+	ValidatorByConsensusAddress   collections.Map[sdk.ConsAddress, sdk.ValAddress]
+	UnbondingType                 collections.Map[uint64, uint64]
+	UnbondingDelegationByValIndex collections.Map[collections.Pair[sdk.ValAddress, sdk.AccAddress], []byte]
 }
 
 // NewKeeper creates a new staking Keeper instance
@@ -98,6 +99,12 @@ func NewKeeper(
 			collcodec.KeyToValueCodec(sdk.ValAddressKey),
 		),
 		UnbondingType: collections.NewMap(sb, types.UnbondingTypeKey, "unbonding_type", collections.Uint64Key, collections.Uint64Value),
+		UnbondingDelegationByValIndex: collections.NewMap(
+			sb, types.UnbondingDelegationByValIndexKey,
+			"unbonding_delegation_by_val_index",
+			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/types/keys.go

@@ -38,12 +38,12 @@ var (
 	ValidatorsByConsAddrKey   = collections.NewPrefix(34) // prefix for each key to a validator index, by pubkey
 	ValidatorsByPowerIndexKey = []byte{0x23}              // prefix for each key to a validator index, sorted by power
 
-	DelegationKey                    = []byte{0x31} // key for a delegation
-	UnbondingDelegationKey           = []byte{0x32} // key for an unbonding-delegation
-	UnbondingDelegationByValIndexKey = []byte{0x33} // prefix for each key for an unbonding-delegation, by validator operator
-	RedelegationKey                  = []byte{0x34} // key for a redelegation
-	RedelegationByValSrcIndexKey     = []byte{0x35} // 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
+	DelegationKey                    = []byte{0x31}              // key for a delegation
+	UnbondingDelegationKey           = []byte{0x32}              // key for an unbonding-delegation
+	UnbondingDelegationByValIndexKey = collections.NewPrefix(51) // prefix for each key for an unbonding-delegation, by validator operator
+	RedelegationKey                  = []byte{0x34}              // key for a redelegation
+	RedelegationByValSrcIndexKey     = []byte{0x35}              // 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
 
 	UnbondingIDKey    = collections.NewPrefix(55) // key for the counter for the incrementing id for UnbondingOperations
 	UnbondingIndexKey = []byte{0x38}              // prefix for an index for looking up unbonding operations by their IDs
@@ -216,7 +216,7 @@ func GetUBDKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
 // GetUBDByValIndexKey creates the index-key for an unbonding delegation, stored by validator-index
 // VALUE: none (key rearrangement used)
 func GetUBDByValIndexKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
-	return append(GetUBDsByValIndexKey(valAddr), address.MustLengthPrefix(delAddr)...)
+	return append(append(UnbondingDelegationByValIndexKey, address.MustLengthPrefix(valAddr)...), address.MustLengthPrefix(delAddr)...)
 }
 
 // GetUBDKeyFromValIndexKey rearranges the ValIndexKey to get the UBDKey
@@ -238,11 +238,6 @@ func GetUBDsKey(delAddr sdk.AccAddress) []byte {
 	return append(UnbondingDelegationKey, address.MustLengthPrefix(delAddr)...)
 }
 
-// GetUBDsByValIndexKey creates the prefix keyspace for the indexes of unbonding delegations for a validator
-func GetUBDsByValIndexKey(valAddr sdk.ValAddress) []byte {
-	return append(UnbondingDelegationByValIndexKey, address.MustLengthPrefix(valAddr)...)
-}
-
 // GetUnbondingDelegationTimeKey creates the prefix for all unbonding delegations from a delegator
 func GetUnbondingDelegationTimeKey(timestamp time.Time) []byte {
 	bz := sdk.FormatTimeBytes(timestamp)