refactor: migrate LastTotalPower to use collections in x/staking (#17026)

CHANGELOG.md

@@ -60,6 +60,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (x/staking) [#17026](https://github.com/cosmos/cosmos-sdk/pull/17026) Use collections for `LastTotalPower`:
+    * remove `Keeper`: `SetLastTotalPower`, `GetLastTotalPower`
 * (staking) [#16959](https://github.com/cosmos/cosmos-sdk/pull/16959) Add validator and consensus address codec as staking keeper arguments.
 * (types) [#16272](https://github.com/cosmos/cosmos-sdk/pull/16272) From now the `FeeGranter` in the `FeeTx` interface takes the byte type instead of string. 
 * (testutil) [#16899](https://github.com/cosmos/cosmos-sdk/pull/16899) The *cli testutil* `QueryBalancesExec` has been removed. Use the gRPC or REST query instead.

x/staking/keeper/genesis.go

@@ -34,7 +34,7 @@ func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) (res
 		panic(err)
 	}
 
-	if err := k.SetLastTotalPower(ctx, data.LastTotalPower); err != nil {
+	if err := k.LastTotalPower.Set(ctx, data.LastTotalPower); err != nil {
 		panic(err)
 	}
 
@@ -238,7 +238,7 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 		panic(err)
 	}
 
-	totalPower, err := k.GetLastTotalPower(ctx)
+	totalPower, err := k.LastTotalPower.Get(ctx)
 	if err != nil {
 		panic(err)
 	}

x/staking/keeper/keeper.go

@@ -6,6 +6,7 @@ import (
 
 	abci "github.com/cometbft/cometbft/abci/types"
 
+	"cosmossdk.io/collections"
 	addresscodec "cosmossdk.io/core/address"
 	storetypes "cosmossdk.io/core/store"
 	"cosmossdk.io/log"
@@ -32,6 +33,9 @@ type Keeper struct {
 	authority             string
 	validatorAddressCodec addresscodec.Codec
 	consensusAddressCodec addresscodec.Codec
+
+	Schema         collections.Schema
+	LastTotalPower collections.Item[math.Int]
 }
 
 // NewKeeper creates a new staking Keeper instance
@@ -44,6 +48,7 @@ func NewKeeper(
 	validatorAddressCodec addresscodec.Codec,
 	consensusAddressCodec addresscodec.Codec,
 ) *Keeper {
+	sb := collections.NewSchemaBuilder(storeService)
 	// ensure bonded and not bonded module accounts are set
 	if addr := ak.GetModuleAddress(types.BondedPoolName); addr == nil {
 		panic(fmt.Sprintf("%s module account has not been set", types.BondedPoolName))
@@ -62,7 +67,7 @@ func NewKeeper(
 		panic("validator and/or consensus address codec are nil")
 	}
 
-	return &Keeper{
+	k := &Keeper{
 		storeService:          storeService,
 		cdc:                   cdc,
 		authKeeper:            ak,
@@ -71,7 +76,15 @@ func NewKeeper(
 		authority:             authority,
 		validatorAddressCodec: validatorAddressCodec,
 		consensusAddressCodec: consensusAddressCodec,
+		LastTotalPower:        collections.NewItem(sb, types.LastTotalPowerKey, "last_total_power", sdk.IntValue),
 	}
+
+	schema, err := sb.Build()
+	if err != nil {
+		panic(err)
+	}
+	k.Schema = schema
+	return k
 }
 
 // Logger returns a module-specific logger.
@@ -100,36 +113,6 @@ func (k *Keeper) SetHooks(sh types.StakingHooks) {
 	k.hooks = sh
 }
 
-// GetLastTotalPower loads the last total validator power.
-func (k Keeper) GetLastTotalPower(ctx context.Context) (math.Int, error) {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := store.Get(types.LastTotalPowerKey)
-	if err != nil {
-		return math.ZeroInt(), err
-	}
-
-	if bz == nil {
-		return math.ZeroInt(), nil
-	}
-
-	var power math.Int
-	if err = power.Unmarshal(bz); err != nil {
-		return math.ZeroInt(), err
-	}
-
-	return power, nil
-}
-
-// SetLastTotalPower sets the last total validator power.
-func (k Keeper) SetLastTotalPower(ctx context.Context, power math.Int) error {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := power.Marshal()
-	if err != nil {
-		return err
-	}
-	return store.Set(types.LastTotalPowerKey, bz)
-}
-
 // GetAuthority returns the x/staking module's authority.
 func (k Keeper) GetAuthority() string {
 	return k.authority

x/staking/keeper/keeper_test.go

@@ -103,8 +103,8 @@ func (s *KeeperTestSuite) TestLastTotalPower() {
 	require := s.Require()
 
 	expTotalPower := math.NewInt(10 ^ 9)
-	require.NoError(keeper.SetLastTotalPower(ctx, expTotalPower))
-	resTotalPower, err := keeper.GetLastTotalPower(ctx)
+	require.NoError(keeper.LastTotalPower.Set(ctx, expTotalPower))
+	resTotalPower, err := keeper.LastTotalPower.Get(ctx)
 	require.NoError(err)
 	require.True(expTotalPower.Equal(resTotalPower))
 }

x/staking/keeper/val_state_change.go

@@ -251,7 +251,7 @@ func (k Keeper) ApplyAndReturnValidatorSetUpdates(ctx context.Context) (updates
 
 	// set total power on lookup index if there are any updates
 	if len(updates) > 0 {
-		if err = k.SetLastTotalPower(ctx, totalPower); err != nil {
+		if err = k.LastTotalPower.Set(ctx, totalPower); err != nil {
 			return nil, err
 		}
 	}

x/staking/types/keys.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"time"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -30,8 +31,8 @@ const (
 var (
 	// Keys for store prefixes
 	// Last* values are constant during a block.
-	LastValidatorPowerKey = []byte{0x11} // prefix for each key to a validator index, for bonded validators
-	LastTotalPowerKey     = []byte{0x12} // prefix for the total power
+	LastValidatorPowerKey = []byte{0x11}              // prefix for each key to a validator index, for bonded validators
+	LastTotalPowerKey     = collections.NewPrefix(18) // prefix for the total power
 
 	ValidatorsKey             = []byte{0x21} // prefix for each key to a validator
 	ValidatorsByConsAddrKey   = []byte{0x22} // prefix for each key to a validator index, by pubkey