refactor(x/staking): Migrate UnbondingIndex to use collections (#17288)

CHANGELOG.md

@@ -63,6 +63,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (x/staking) [#17288](https://github.com/cosmos/cosmos-sdk/pull/17288) Use collections for `UnbondingIndex`:
+    * remove from `types`: `GetUnbondingIndexKey`.
 * (x/staking) [#17256](https://github.com/cosmos/cosmos-sdk/pull/17256) Use collections for `UnbondingID`.
 * (x/staking) [#17260](https://github.com/cosmos/cosmos-sdk/pull/17260) Use collections for `ValidatorByConsAddr`:
     * remove from `types`: `GetValidatorByConsAddrKey`

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]
+	UnbondingIndex              collections.Map[uint64, []byte]
 }
 
 // NewKeeper creates a new staking Keeper instance
@@ -97,7 +98,8 @@ 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),
+		UnbondingIndex: collections.NewMap(sb, types.UnbondingIndexKey, "unbonding_index", collections.Uint64Key, collections.BytesValue),
 	}
 
 	schema, err := sb.Build()

x/staking/keeper/unbonding.go

@@ -24,8 +24,7 @@ func (k Keeper) IncrementUnbondingID(ctx context.Context) (unbondingID uint64, e
 
 // DeleteUnbondingIndex removes a mapping from UnbondingId to unbonding operation
 func (k Keeper) DeleteUnbondingIndex(ctx context.Context, id uint64) error {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Delete(types.GetUnbondingIndexKey(id))
+	return k.UnbondingIndex.Remove(ctx, id)
 }
 
 // GetUnbondingType returns the enum type of unbonding which is any of
@@ -48,8 +47,11 @@ func (k Keeper) SetUnbondingType(ctx context.Context, id uint64, unbondingType t
 func (k Keeper) GetUnbondingDelegationByUnbondingID(ctx context.Context, id uint64) (ubd types.UnbondingDelegation, err error) {
 	store := k.storeService.OpenKVStore(ctx)
 
-	ubdKey, err := store.Get(types.GetUnbondingIndexKey(id))
+	ubdKey, err := k.UnbondingIndex.Get(ctx, id)
 	if err != nil {
+		if errors.Is(err, collections.ErrNotFound) {
+			return types.UnbondingDelegation{}, types.ErrNoUnbondingDelegation
+		}
 		return types.UnbondingDelegation{}, err
 	}
 
@@ -79,8 +81,11 @@ func (k Keeper) GetUnbondingDelegationByUnbondingID(ctx context.Context, id uint
 func (k Keeper) GetRedelegationByUnbondingID(ctx context.Context, id uint64) (red types.Redelegation, err error) {
 	store := k.storeService.OpenKVStore(ctx)
 
-	redKey, err := store.Get(types.GetUnbondingIndexKey(id))
+	redKey, err := k.UnbondingIndex.Get(ctx, id)
 	if err != nil {
+		if errors.Is(err, collections.ErrNotFound) {
+			return types.Redelegation{}, types.ErrNoRedelegation
+		}
 		return types.Redelegation{}, err
 	}
 
@@ -110,8 +115,11 @@ func (k Keeper) GetRedelegationByUnbondingID(ctx context.Context, id uint64) (re
 func (k Keeper) GetValidatorByUnbondingID(ctx context.Context, id uint64) (val types.Validator, err error) {
 	store := k.storeService.OpenKVStore(ctx)
 
-	valKey, err := store.Get(types.GetUnbondingIndexKey(id))
+	valKey, err := k.UnbondingIndex.Get(ctx, id)
 	if err != nil {
+		if errors.Is(err, collections.ErrNotFound) {
+			return types.Validator{}, types.ErrNoValidatorFound
+		}
 		return types.Validator{}, err
 	}
 
@@ -141,7 +149,6 @@ func (k Keeper) GetValidatorByUnbondingID(ctx context.Context, id uint64) (val t
 // by the unbondingID of an UnbondingDelegationEntry that it contains Note, it does not
 // set the unbonding delegation itself, use SetUnbondingDelegation(ctx, ubd) for that
 func (k Keeper) SetUnbondingDelegationByUnbondingID(ctx context.Context, ubd types.UnbondingDelegation, id uint64) error {
-	store := k.storeService.OpenKVStore(ctx)
 	delAddr, err := k.authKeeper.AddressCodec().StringToBytes(ubd.DelegatorAddress)
 	if err != nil {
 		return err
@@ -152,7 +159,7 @@ func (k Keeper) SetUnbondingDelegationByUnbondingID(ctx context.Context, ubd typ
 	}
 
 	ubdKey := types.GetUBDKey(delAddr, valAddr)
-	if err = store.Set(types.GetUnbondingIndexKey(id), ubdKey); err != nil {
+	if err = k.UnbondingIndex.Set(ctx, id, ubdKey); err != nil {
 		return err
 	}
 
@@ -163,8 +170,6 @@ func (k Keeper) SetUnbondingDelegationByUnbondingID(ctx context.Context, ubd typ
 // SetRedelegationByUnbondingID sets an index to look up an Redelegation by the unbondingID of an RedelegationEntry that it contains
 // Note, it does not set the redelegation itself, use SetRedelegation(ctx, red) for that
 func (k Keeper) SetRedelegationByUnbondingID(ctx context.Context, red types.Redelegation, id uint64) error {
-	store := k.storeService.OpenKVStore(ctx)
-
 	delAddr, err := k.authKeeper.AddressCodec().StringToBytes(red.DelegatorAddress)
 	if err != nil {
 		return err
@@ -181,7 +186,7 @@ func (k Keeper) SetRedelegationByUnbondingID(ctx context.Context, red types.Rede
 	}
 
 	redKey := types.GetREDKey(delAddr, valSrcAddr, valDstAddr)
-	if err = store.Set(types.GetUnbondingIndexKey(id), redKey); err != nil {
+	if err = k.UnbondingIndex.Set(ctx, id, redKey); err != nil {
 		return err
 	}
 
@@ -192,15 +197,13 @@ func (k Keeper) SetRedelegationByUnbondingID(ctx context.Context, red types.Rede
 // SetValidatorByUnbondingID sets an index to look up a Validator by the unbondingID corresponding to its current unbonding
 // Note, it does not set the validator itself, use SetValidator(ctx, val) for that
 func (k Keeper) SetValidatorByUnbondingID(ctx context.Context, val types.Validator, id uint64) error {
-	store := k.storeService.OpenKVStore(ctx)
-
 	valAddr, err := k.validatorAddressCodec.StringToBytes(val.OperatorAddress)
 	if err != nil {
 		return err
 	}
 
 	valKey := types.GetValidatorKey(valAddr)
-	if err = store.Set(types.GetUnbondingIndexKey(id), valKey); err != nil {
+	if err = k.UnbondingIndex.Set(ctx, id, valKey); err != nil {
 		return err
 	}
 

x/staking/types/keys.go

@@ -46,7 +46,7 @@ var (
 	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
+	UnbondingIndexKey = collections.NewPrefix(56) // prefix for an index for looking up unbonding operations by their IDs
 	UnbondingTypeKey  = collections.NewPrefix(57) // prefix for an index containing the type of unbonding operations
 
 	UnbondingQueueKey    = []byte{0x41} // prefix for the timestamps in unbonding queue
@@ -71,13 +71,6 @@ const (
 	UnbondingType_ValidatorUnbonding
 )
 
-// GetUnbondingIndexKey returns a key for the index for looking up UnbondingDelegations by the UnbondingDelegationEntries they contain
-func GetUnbondingIndexKey(id uint64) []byte {
-	bz := make([]byte, 8)
-	binary.BigEndian.PutUint64(bz, id)
-	return append(UnbondingIndexKey, bz...)
-}
-
 // GetValidatorKey creates the key for the validator with address
 // VALUE: staking/Validator
 func GetValidatorKey(operatorAddr sdk.ValAddress) []byte {