Revert iterator optimization

x/nft/keeper/class.go

@@ -3,8 +3,6 @@ package keeper
 import (
 	"context"
 
-	"github.com/cosmos/cosmos-sdk/codec"
-
 	"cosmossdk.io/errors"
 	storetypes "cosmossdk.io/store/types"
 	"cosmossdk.io/x/nft"
@@ -68,18 +66,6 @@ func (k Keeper) GetClasses(ctx context.Context) (classes []*nft.Class) {
 	return
 }
 
-// RawIterateClasses iterate on the raw classes store date
-func (k Keeper) RawIterateClasses(ctx context.Context, cb func(cdc codec.BinaryCodec, key, value []byte) bool) {
-	store := k.KVStoreService.OpenKVStore(ctx)
-	iterator := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), ClassKey)
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		if cb(k.cdc, iterator.Key(), iterator.Value()) {
-			return
-		}
-	}
-}
-
 // HasClass determines whether the specified classID exist
 func (k Keeper) HasClass(ctx context.Context, classID string) bool {
 	store := k.KVStoreService.OpenKVStore(ctx)

x/nft/simulation/operations.go

@@ -151,8 +151,8 @@ func randNFT(ctx sdk.Context, r *rand.Rand, k keeper.Keeper, minter sdk.AccAddre
 
 // randClass picks a random Class.
 func randClass(ctx sdk.Context, r *rand.Rand, k keeper.Keeper) (nft.Class, error) {
-	classCount := totalClasses(ctx, k)
-	if classCount == 0 {
+	classes := k.GetClasses(ctx)
+	if len(classes) == 0 {
 		c := nft.Class{
 			Id:          simtypes.RandStringOfLength(r, 10),
 			Name:        simtypes.RandStringOfLength(r, 10),
@@ -166,34 +166,5 @@ func randClass(ctx sdk.Context, r *rand.Rand, k keeper.Keeper) (nft.Class, error
 		}
 		return c, nil
 	}
-	return loadClassN(ctx, k, r.Intn(classCount)), nil
-}
-
-func totalClasses(ctx sdk.Context, k keeper.Keeper) int {
-	var count int
-	k.RawIterateClasses(ctx, func(_ codec.BinaryCodec, _, _ []byte) bool {
-		count++
-		return false
-	})
-	return count
-}
-
-func loadClassN(ctx sdk.Context, k keeper.Keeper, atPos int) nft.Class {
-	var (
-		currentPos int
-		result     *nft.Class
-	)
-	k.RawIterateClasses(ctx, func(cdc codec.BinaryCodec, _, value []byte) bool {
-		if currentPos == atPos {
-			result = new(nft.Class)
-			cdc.MustUnmarshal(value, result)
-			return true
-		}
-		currentPos++
-		return false
-	})
-	if result == nil {
-		panic("not found")
-	}
-	return *result
+	return *classes[r.Intn(len(classes))], nil
 }

x/staking/keeper/validator.go

@@ -7,9 +7,6 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/cosmos/cosmos-sdk/codec"
-	"github.com/cosmos/cosmos-sdk/runtime"
-
 	gogotypes "github.com/cosmos/gogoproto/types"
 
 	"cosmossdk.io/collections"
@@ -282,18 +279,6 @@ func (k Keeper) GetAllValidators(ctx context.Context) (validators []types.Valida
 	return validators, nil
 }
 
-// RawIterateValidators iterate on the raw validator store date
-func (k Keeper) RawIterateValidators(ctx context.Context, cb func(cdc codec.BinaryCodec, key, value []byte) bool) {
-	store := k.KVStoreService.OpenKVStore(ctx)
-	iterator := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), types.ValidatorsKey)
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		if cb(k.cdc, iterator.Key(), iterator.Value()) {
-			return
-		}
-	}
-}
-
 // GetValidators returns a given amount of all the validators
 func (k Keeper) GetValidators(ctx context.Context, maxRetrieve uint32) (validators []types.Validator, err error) {
 	store := k.KVStoreService.OpenKVStore(ctx)

x/staking/simulation/operations.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/testutil"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	simtypes "github.com/cosmos/cosmos-sdk/types/simulation"
 	"github.com/cosmos/cosmos-sdk/x/simulation"
@@ -234,7 +235,17 @@ func SimulateMsgEditValidator(
 		r *rand.Rand, app simtypes.AppEntrypoint, ctx sdk.Context, accs []simtypes.Account, chainID string,
 	) (simtypes.OperationMsg, []simtypes.FutureOperation, error) {
 		msgType := sdk.MsgTypeURL(&types.MsgEditValidator{})
-		val, ok := randValidator(ctx, r, k)
+
+		vals, err := k.GetAllValidators(ctx)
+		if err != nil {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to get validators"), nil, err
+		}
+
+		if len(vals) == 0 {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "number of validators equal zero"), nil, nil
+		}
+
+		val, ok := testutil.RandSliceElem(r, vals)
 		if !ok {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to pick a validator"), nil, nil
 		}
@@ -304,11 +315,20 @@ func SimulateMsgDelegate(
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "bond denom not found"), nil, err
 		}
 
-		val, ok := randValidator(ctx, r, k)
+		vals, err := k.GetAllValidators(ctx)
+		if err != nil {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to get validators"), nil, err
+		}
+
+		if len(vals) == 0 {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "number of validators equal zero"), nil, nil
+		}
+
+		simAccount, _ := simtypes.RandomAcc(r, accs)
+		val, ok := testutil.RandSliceElem(r, vals)
 		if !ok {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to pick a validator"), nil, nil
 		}
-		simAccount, _ := simtypes.RandomAcc(r, accs)
 
 		if val.InvalidExRate() {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "validator's invalid echange rate"), nil, nil
@@ -373,7 +393,16 @@ func SimulateMsgUndelegate(
 	) (simtypes.OperationMsg, []simtypes.FutureOperation, error) {
 		msgType := sdk.MsgTypeURL(&types.MsgUndelegate{})
 
-		val, ok := randValidator(ctx, r, k)
+		vals, err := k.GetAllValidators(ctx)
+		if err != nil {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to get validators"), nil, err
+		}
+
+		if len(vals) == 0 {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "number of validators equal zero"), nil, nil
+		}
+
+		val, ok := testutil.RandSliceElem(r, vals)
 		if !ok {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "validator is not ok"), nil, nil
 		}
@@ -483,11 +512,20 @@ func SimulateMsgCancelUnbondingDelegate(
 	) (simtypes.OperationMsg, []simtypes.FutureOperation, error) {
 		msgType := sdk.MsgTypeURL(&types.MsgCancelUnbondingDelegation{})
 
-		val, ok := randValidator(ctx, r, k)
+		vals, err := k.GetAllValidators(ctx)
+		if err != nil {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to get validators"), nil, err
+		}
+
+		if len(vals) == 0 {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "number of validators equal zero"), nil, nil
+		}
+
+		simAccount, _ := simtypes.RandomAcc(r, accs)
+		val, ok := testutil.RandSliceElem(r, vals)
 		if !ok {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "validator is not ok"), nil, nil
 		}
-		simAccount, _ := simtypes.RandomAcc(r, accs)
 
 		if val.IsJailed() || val.InvalidExRate() {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "validator is jailed"), nil, nil
@@ -578,7 +616,16 @@ func SimulateMsgBeginRedelegate(
 	) (simtypes.OperationMsg, []simtypes.FutureOperation, error) {
 		msgType := sdk.MsgTypeURL(&types.MsgBeginRedelegate{})
 
-		srcVal, ok := randValidator(ctx, r, k)
+		allVals, err := k.GetAllValidators(ctx)
+		if err != nil {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to get validators"), nil, err
+		}
+
+		if len(allVals) == 0 {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "number of validators equal zero"), nil, nil
+		}
+
+		srcVal, ok := testutil.RandSliceElem(r, allVals)
 		if !ok {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to pick validator"), nil, nil
 		}
@@ -615,7 +662,7 @@ func SimulateMsgBeginRedelegate(
 		}
 
 		// get random destination validator
-		destVal, ok := randValidator(ctx, r, k)
+		destVal, ok := testutil.RandSliceElem(r, allVals)
 		if !ok {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to pick validator"), nil, nil
 		}
@@ -713,7 +760,16 @@ func SimulateMsgRotateConsPubKey(txGen client.TxConfig, ak types.AccountKeeper,
 	) (simtypes.OperationMsg, []simtypes.FutureOperation, error) {
 		msgType := sdk.MsgTypeURL(&types.MsgRotateConsPubKey{})
 
-		val, ok := randValidator(ctx, r, k)
+		vals, err := k.GetAllValidators(ctx)
+		if err != nil {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to get validators"), nil, err
+		}
+
+		if len(vals) == 0 {
+			return simtypes.NoOpMsg(types.ModuleName, msgType, "number of validators equal zero"), nil, nil
+		}
+
+		val, ok := testutil.RandSliceElem(r, vals)
 		if !ok {
 			return simtypes.NoOpMsg(types.ModuleName, msgType, "unable to pick a validator"), nil, nil
 		}
@@ -815,40 +871,3 @@ func SimulateMsgRotateConsPubKey(txGen client.TxConfig, ak types.AccountKeeper,
 		return simulation.GenAndDeliverTxWithRandFees(txCtx)
 	}
 }
-
-func randValidator(ctx sdk.Context, r *rand.Rand, k *keeper.Keeper) (types.Validator, bool) {
-	valCount := totalValidators(ctx, k)
-	if valCount == 0 {
-		return types.Validator{}, false
-	}
-	return loadValidatorN(ctx, k, r.Intn(valCount)), true
-}
-
-func totalValidators(ctx sdk.Context, k *keeper.Keeper) int {
-	var count int
-	k.RawIterateValidators(ctx, func(_ codec.BinaryCodec, _, _ []byte) bool {
-		count++
-		return false
-	})
-	return count
-}
-
-func loadValidatorN(ctx sdk.Context, k *keeper.Keeper, atPos int) types.Validator {
-	var (
-		currentPos int
-		result     *types.Validator
-	)
-	k.RawIterateValidators(ctx, func(cdc codec.BinaryCodec, _, value []byte) bool {
-		if currentPos == atPos {
-			result = new(types.Validator)
-			cdc.MustUnmarshal(value, result)
-			return true
-		}
-		currentPos++
-		return false
-	})
-	if result == nil {
-		panic("not found")
-	}
-	return *result
-}