refactor(x/staking): Migrate LastValidatorPower to Collections (#17498)

CHANGELOG.md

@@ -59,6 +59,9 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (x/staking) [#17498](https://github.com/cosmos/cosmos-sdk/pull/17498) Use collections for `LastValidatorPower`:
+    * remove from `types`: `GetLastValidatorPowerKey`
+    * remove from `Keeper`: `LastValidatorsIterator`, `IterateLastValidators`
 * (x/staking) [#17291](https://github.com/cosmos/cosmos-sdk/pull/17291) Use collections for `UnbondingDelegationByValIndex`:
     * remove from `types`: `GetUBDKeyFromValIndexKey`, `GetUBDsByValIndexKey`, `GetUBDByValIndexKey`
 * (x/slashing) [#17568](https://github.com/cosmos/cosmos-sdk/pull/17568) Use collections for `ValidatorMissedBlockBitmap`:

x/staking/genesis.go

@@ -13,16 +13,21 @@ import (
 
 // WriteValidators returns a slice of bonded genesis validators.
 func WriteValidators(ctx sdk.Context, keeper *keeper.Keeper) (vals []cmttypes.GenesisValidator, returnErr error) {
-	err := keeper.IterateLastValidators(ctx, func(_ int64, validator types.ValidatorI) (stop bool) {
+	err := keeper.LastValidatorPower.Walk(ctx, nil, func(key, value []byte) (bool, error) {
+		validator, err := keeper.GetValidator(ctx, key)
+		if err != nil {
+			return true, err
+		}
+
 		pk, err := validator.ConsPubKey()
 		if err != nil {
 			returnErr = err
-			return true
+			return true, err
 		}
 		cmtPk, err := cryptocodec.ToCmtPubKeyInterface(pk)
 		if err != nil {
 			returnErr = err
-			return true
+			return true, err
 		}
 
 		vals = append(vals, cmttypes.GenesisValidator{
@@ -32,13 +37,13 @@ func WriteValidators(ctx sdk.Context, keeper *keeper.Keeper) (vals []cmttypes.Ge
 			Name:    validator.GetMoniker(),
 		})
 
-		return false
+		return false, nil
 	})
 	if err != nil {
 		return nil, err
 	}
 
-	return
+	return vals, returnErr
 }
 
 // ValidateGenesis validates the provided staking genesis state to ensure the

x/staking/keeper/alias_functions.go

@@ -70,33 +70,6 @@ func (k Keeper) IterateBondedValidatorsByPower(ctx context.Context, fn func(inde
 	return nil
 }
 
-// IterateLastValidators iterates through the active validator set and perform the provided function
-func (k Keeper) IterateLastValidators(ctx context.Context, fn func(index int64, validator types.ValidatorI) (stop bool)) error {
-	iterator, err := k.LastValidatorsIterator(ctx)
-	if err != nil {
-		return err
-	}
-	defer iterator.Close()
-
-	i := int64(0)
-
-	for ; iterator.Valid(); iterator.Next() {
-		address := types.AddressFromLastValidatorPowerKey(iterator.Key())
-
-		validator, err := k.GetValidator(ctx, address)
-		if err != nil {
-			return err
-		}
-
-		stop := fn(i, validator) // XXX is this safe will the validator unexposed fields be able to get written to?
-		if stop {
-			break
-		}
-		i++
-	}
-	return nil
-}
-
 // Validator gets the Validator interface for a particular address
 func (k Keeper) Validator(ctx context.Context, address sdk.ValAddress) (types.ValidatorI, error) {
 	return k.GetValidator(ctx, address)

x/staking/keeper/keeper.go

@@ -51,6 +51,7 @@ type Keeper struct {
 	RedelegationsByValDst         collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
 	RedelegationsByValSrc         collections.Map[collections.Triple[[]byte, []byte, []byte], []byte]
 	UnbondingDelegationByValIndex collections.Map[collections.Pair[[]byte, []byte], []byte]
+	LastValidatorPower            collections.Map[[]byte, []byte]
 }
 
 // NewKeeper creates a new staking Keeper instance
@@ -146,6 +147,7 @@ func NewKeeper(
 			),
 			collections.BytesValue,
 		),
+		LastValidatorPower: collections.NewMap(sb, types.LastValidatorPowerKey, "last_validator_power", sdk.LengthPrefixedBytesKey, collections.BytesValue), // sdk.LengthPrefixedBytesKey is needed to retain state compatibility
 		// key format is: 54 | lengthPrefixedBytes(DstValAddr) | lengthPrefixedBytes(AccAddr) | lengthPrefixedBytes(SrcValAddr)
 		RedelegationsByValDst: collections.NewMap(
 			sb, types.RedelegationByValDstIndexKey,

x/staking/keeper/keeper_test.go

@@ -6,6 +6,7 @@ import (
 
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	cmttime "github.com/cometbft/cometbft/types/time"
+	gogotypes "github.com/cosmos/gogoproto/types"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/suite"
 
@@ -195,6 +196,47 @@ func getValidatorKey(operatorAddr sdk.ValAddress) []byte {
 	return append(validatorsKey, addresstypes.MustLengthPrefix(operatorAddr)...)
 }
 
+// getLastValidatorPowerKey creates the bonded validator index key for an operator address
+func getLastValidatorPowerKey(operator sdk.ValAddress) []byte {
+	lastValidatorPowerKey := []byte{0x11}
+	return append(lastValidatorPowerKey, addresstypes.MustLengthPrefix(operator)...)
+}
+
+func (s *KeeperTestSuite) TestLastTotalPowerMigrationToColls() {
+	s.SetupTest()
+
+	_, valAddrs := createValAddrs(100)
+
+	err := testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			bz, err := s.cdc.Marshal(&gogotypes.Int64Value{Value: i})
+			s.Require().NoError(err)
+
+			s.ctx.KVStore(s.key).Set(getLastValidatorPowerKey(valAddrs[i]), bz)
+		},
+		"f28811f2b0a0ab9db60cdcae93680faff9dbadd4a3a8a2d088bb19b0428ad3a9",
+	)
+	s.Require().NoError(err)
+
+	err = testutil.DiffCollectionsMigration(
+		s.ctx,
+		s.key,
+		100,
+		func(i int64) {
+			bz, err := s.cdc.Marshal(&gogotypes.Int64Value{Value: i})
+			s.Require().NoError(err)
+
+			err = s.stakingKeeper.LastValidatorPower.Set(s.ctx, valAddrs[i], bz)
+			s.Require().NoError(err)
+		},
+		"f28811f2b0a0ab9db60cdcae93680faff9dbadd4a3a8a2d088bb19b0428ad3a9",
+	)
+	s.Require().NoError(err)
+}
+
 func (s *KeeperTestSuite) TestSrcRedelegationsMigrationToColls() {
 	s.SetupTest()
 

x/staking/keeper/val_state_change.go

@@ -463,23 +463,17 @@ type validatorsByAddr map[string][]byte
 func (k Keeper) getLastValidatorsByAddr(ctx context.Context) (validatorsByAddr, error) {
 	last := make(validatorsByAddr)
 
-	iterator, err := k.LastValidatorsIterator(ctx)
-	if err != nil {
-		return nil, err
-	}
-	defer iterator.Close()
-
-	for ; iterator.Valid(); iterator.Next() {
-		// extract the validator address from the key (prefix is 1-byte, addrLen is 1-byte)
-		valAddr := types.AddressFromLastValidatorPowerKey(iterator.Key())
-		valAddrStr, err := k.validatorAddressCodec.BytesToString(valAddr)
+	err := k.LastValidatorPower.Walk(ctx, nil, func(key, value []byte) (bool, error) {
+		valAddrStr, err := k.validatorAddressCodec.BytesToString(key)
 		if err != nil {
-			return nil, err
+			return true, err
 		}
 
-		powerBytes := iterator.Value()
-		last[valAddrStr] = make([]byte, len(powerBytes))
-		copy(last[valAddrStr], powerBytes)
+		last[valAddrStr] = value
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	return last, nil

x/staking/keeper/validator.go

@@ -337,8 +337,7 @@ func (k Keeper) ValidatorsPowerStoreIterator(ctx context.Context) (corestore.Ite
 // GetLastValidatorPower loads the last validator power.
 // Returns zero if the operator was not a validator last block.
 func (k Keeper) GetLastValidatorPower(ctx context.Context, operator sdk.ValAddress) (power int64, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := store.Get(types.GetLastValidatorPowerKey(operator))
+	bz, err := k.LastValidatorPower.Get(ctx, operator)
 	if err != nil {
 		return 0, err
 	}
@@ -358,81 +357,67 @@ func (k Keeper) GetLastValidatorPower(ctx context.Context, operator sdk.ValAddre
 
 // SetLastValidatorPower sets the last validator power.
 func (k Keeper) SetLastValidatorPower(ctx context.Context, operator sdk.ValAddress, power int64) error {
-	store := k.storeService.OpenKVStore(ctx)
 	bz, err := k.cdc.Marshal(&gogotypes.Int64Value{Value: power})
 	if err != nil {
 		return err
 	}
-	return store.Set(types.GetLastValidatorPowerKey(operator), bz)
+	return k.LastValidatorPower.Set(ctx, operator, bz)
 }
 
 // DeleteLastValidatorPower deletes the last validator power.
 func (k Keeper) DeleteLastValidatorPower(ctx context.Context, operator sdk.ValAddress) error {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Delete(types.GetLastValidatorPowerKey(operator))
-}
-
-// lastValidatorsIterator returns an iterator for the consensus validators in the last block
-func (k Keeper) LastValidatorsIterator(ctx context.Context) (corestore.Iterator, error) {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Iterator(types.LastValidatorPowerKey, storetypes.PrefixEndBytes(types.LastValidatorPowerKey))
+	return k.LastValidatorPower.Remove(ctx, operator)
 }
 
 // IterateLastValidatorPowers iterates over last validator powers.
 func (k Keeper) IterateLastValidatorPowers(ctx context.Context, handler func(operator sdk.ValAddress, power int64) (stop bool)) error {
-	iter, err := k.LastValidatorsIterator(ctx)
-	if err != nil {
-		return err
-	}
-
-	for ; iter.Valid(); iter.Next() {
-		addr := sdk.ValAddress(types.AddressFromLastValidatorPowerKey(iter.Key()))
+	err := k.LastValidatorPower.Walk(ctx, nil, func(key, value []byte) (bool, error) {
+		addr := sdk.ValAddress(key)
 		intV := &gogotypes.Int64Value{}
 
-		if err = k.cdc.Unmarshal(iter.Value(), intV); err != nil {
-			return err
+		if err := k.cdc.Unmarshal(value, intV); err != nil {
+			return true, err
 		}
 
 		if handler(addr, intV.GetValue()) {
-			break
+			return true, nil
 		}
+		return false, nil
+	})
+	if err != nil {
+		return err
 	}
 
 	return nil
 }
 
 // GetLastValidators gets the group of the bonded validators
 func (k Keeper) GetLastValidators(ctx context.Context) (validators []types.Validator, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-
 	// add the actual validator power sorted store
 	maxValidators, err := k.MaxValidators(ctx)
 	if err != nil {
 		return nil, err
 	}
 	validators = make([]types.Validator, maxValidators)
 
-	iterator, err := store.Iterator(types.LastValidatorPowerKey, storetypes.PrefixEndBytes(types.LastValidatorPowerKey))
-	if err != nil {
-		return nil, err
-	}
-	defer iterator.Close()
-
 	i := 0
-	for ; iterator.Valid(); iterator.Next() {
+	err = k.LastValidatorPower.Walk(ctx, nil, func(key, value []byte) (bool, error) {
 		// sanity check
 		if i >= int(maxValidators) {
 			panic("more validators than maxValidators found")
 		}
 
-		address := types.AddressFromLastValidatorPowerKey(iterator.Key())
-		validator, err := k.GetValidator(ctx, address)
+		validator, err := k.GetValidator(ctx, key)
 		if err != nil {
-			return nil, err
+			return true, err
 		}
 
 		validators[i] = validator
 		i++
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	return validators[:i], nil // trim

x/staking/keeper/validator_test.go

@@ -6,6 +6,7 @@ import (
 	abci "github.com/cometbft/cometbft/abci/types"
 	"github.com/golang/mock/gomock"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -82,7 +83,7 @@ func (s *KeeperTestSuite) TestValidator() {
 	require.Equal(power, resPower)
 	require.NoError(keeper.DeleteLastValidatorPower(ctx, valAddr))
 	resPower, err = keeper.GetLastValidatorPower(ctx, valAddr)
-	require.NoError(err)
+	require.Error(err, collections.ErrNotFound)
 	require.Equal(int64(0), resPower)
 }
 

x/staking/migrations/v2/store_test.go

@@ -46,7 +46,7 @@ func TestStoreMigration(t *testing.T) {
 		{
 			"LastValidatorPowerKey",
 			v1.GetLastValidatorPowerKey(valAddr1),
-			types.GetLastValidatorPowerKey(valAddr1),
+			getLastValidatorPowerKey(valAddr1),
 		},
 		{
 			"LastTotalPowerKey",
@@ -141,6 +141,10 @@ func TestStoreMigration(t *testing.T) {
 	}
 }
 
+func getLastValidatorPowerKey(operator sdk.ValAddress) []byte {
+	return append(types.LastValidatorPowerKey, sdkaddress.MustLengthPrefix(operator)...)
+}
+
 func getUBDByValIndexKey(delAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
 	return append(append(types.UnbondingDelegationByValIndexKey, sdkaddress.MustLengthPrefix(valAddr)...), sdkaddress.MustLengthPrefix(delAddr)...)
 }

x/staking/types/expected_keepers.go

@@ -52,10 +52,6 @@ type ValidatorSet interface {
 	IterateBondedValidatorsByPower(context.Context,
 		func(index int64, validator ValidatorI) (stop bool)) error
 
-	// iterate through the consensus validator set of the last block by operator address, execute func for each validator
-	IterateLastValidators(context.Context,
-		func(index int64, validator ValidatorI) (stop bool)) error
-
 	Validator(context.Context, sdk.ValAddress) (ValidatorI, error)            // get a particular validator by operator address
 	ValidatorByConsAddr(context.Context, sdk.ConsAddress) (ValidatorI, error) // get a particular validator by consensus address
 	TotalBondedTokens(context.Context) (math.Int, error)                      // total bonded tokens within the validator set

x/staking/types/keys.go

@@ -32,7 +32,7 @@ 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
+	LastValidatorPowerKey = collections.NewPrefix(17) // prefix for each key to a validator index, for bonded validators
 	LastTotalPowerKey     = collections.NewPrefix(18) // prefix for the total power
 
 	ValidatorsKey             = collections.NewPrefix(33) // prefix for each key to a validator
@@ -128,11 +128,6 @@ func GetValidatorsByPowerIndexKey(validator Validator, powerReduction math.Int,
 	return key
 }
 
-// GetLastValidatorPowerKey creates the bonded validator index key for an operator address
-func GetLastValidatorPowerKey(operator sdk.ValAddress) []byte {
-	return append(LastValidatorPowerKey, address.MustLengthPrefix(operator)...)
-}
-
 // ParseValidatorPowerRankKey parses the validators operator address from power rank key
 func ParseValidatorPowerRankKey(key []byte) (operAddr []byte) {
 	powerBytesLen := 8