From 92dffb5fa87f38c8f312e88601280256bd01df7a Mon Sep 17 00:00:00 2001
From: Likhita Polavarapu <78951027+likhita-809@users.noreply.github.com>
Date: Thu, 3 Aug 2023 18:20:16 +0530
Subject: [PATCH] refactor(x/staking): Migrate UnbondingID to collections
 (#17256)

---
 CHANGELOG.md                  |  1 +
 x/staking/keeper/keeper.go    |  2 ++
 x/staking/keeper/unbonding.go | 17 +----------------
 x/staking/types/keys.go       |  2 +-
 4 files changed, 5 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 78630e85b540..d25ab2a7f81c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -60,6 +60,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (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) [#17248](https://github.com/cosmos/cosmos-sdk/pull/17248) Use collections for `UnbondingType`.
diff --git a/x/staking/keeper/keeper.go b/x/staking/keeper/keeper.go
index 938dfa855ece..351ea8cf8612 100644
--- a/x/staking/keeper/keeper.go
+++ b/x/staking/keeper/keeper.go
@@ -38,6 +38,7 @@ type Keeper struct {
 	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]
 }
@@ -89,6 +90,7 @@ func NewKeeper(
 			collections.PairKeyCodec(sdk.LengthPrefixedAddressKey(sdk.ValAddressKey), sdk.AccAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
 			collections.BytesValue,
 		),
+		UnbondingID: collections.NewSequence(sb, types.UnbondingIDKey, "unbonding_id"),
 		ValidatorByConsensusAddress: collections.NewMap(
 			sb, types.ValidatorsByConsAddrKey,
 			"validator_by_cons_addr",
diff --git a/x/staking/keeper/unbonding.go b/x/staking/keeper/unbonding.go
index 0c457733c869..bc09bf8b488e 100644
--- a/x/staking/keeper/unbonding.go
+++ b/x/staking/keeper/unbonding.go
@@ -2,7 +2,6 @@ package keeper
 
 import (
 	"context"
-	"encoding/binary"
 	"errors"
 
 	"cosmossdk.io/collections"
@@ -14,26 +13,12 @@ import (
 
 // IncrementUnbondingID increments and returns a unique ID for an unbonding operation
 func (k Keeper) IncrementUnbondingID(ctx context.Context) (unbondingID uint64, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := store.Get(types.UnbondingIDKey)
+	unbondingID, err = k.UnbondingID.Next(ctx)
 	if err != nil {
 		return 0, err
 	}
-
-	if bz != nil {
-		unbondingID = binary.BigEndian.Uint64(bz)
-	}
-
 	unbondingID++
 
-	// Convert back into bytes for storage
-	bz = make([]byte, 8)
-	binary.BigEndian.PutUint64(bz, unbondingID)
-
-	if err = store.Set(types.UnbondingIDKey, bz); err != nil {
-		return 0, err
-	}
-
 	return unbondingID, err
 }
 
diff --git a/x/staking/types/keys.go b/x/staking/types/keys.go
index 8572eec7b033..61da6f3eda24 100644
--- a/x/staking/types/keys.go
+++ b/x/staking/types/keys.go
@@ -45,7 +45,7 @@ var (
 	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    = []byte{0x37}              // key for the counter for the incrementing id for UnbondingOperations
+	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
 	UnbondingTypeKey  = collections.NewPrefix(57) // prefix for an index containing the type of unbonding operations