fix: Bank module init genesis optimization (#9428)

* optimize the bank module genesis initialization

* remove k.setBalances & k.clearBalances and update changelog

* fix lint

Co-authored-by: Aaron Craelius <aaron@regen.network>

CHANGELOG.md

@@ -36,6 +36,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ## [Unreleased]
 
+* [\#9428](https://github.com/cosmos/cosmos-sdk/pull/9428) Optimize bank InitGenesis. Removed bank keeper's `k.setBalances` and `k.clearBalances`. Added `k.initBalances`.
 * [\#9231](https://github.com/cosmos/cosmos-sdk/pull/9231) Remove redundant staking errors.
 * [\#9205](https://github.com/cosmos/cosmos-sdk/pull/9205) Improve readability in `abci` handleQueryP2P
 * [\#9235](https://github.com/cosmos/cosmos-sdk/pull/9235) CreateMembershipProof/CreateNonMembershipProof now returns an error

x/bank/keeper/genesis.go

@@ -21,7 +21,7 @@ func (k BaseKeeper) InitGenesis(ctx sdk.Context, genState *types.GenesisState) {
 			panic(err)
 		}
 
-		if err := k.setBalances(ctx, addr, balance.Coins); err != nil {
+		if err := k.initBalances(ctx, addr, balance.Coins); err != nil {
 			panic(fmt.Errorf("error on setting balances %w", err))
 		}
 

x/bank/keeper/send.go

@@ -227,31 +227,20 @@ func (k BaseSendKeeper) addCoins(ctx sdk.Context, addr sdk.AccAddress, amt sdk.C
 	return nil
 }
 
-// clearBalances removes all balances for a given account by address.
-func (k BaseSendKeeper) clearBalances(ctx sdk.Context, addr sdk.AccAddress) {
-	keys := [][]byte{}
-	k.IterateAccountBalances(ctx, addr, func(balance sdk.Coin) bool {
-		keys = append(keys, []byte(balance.Denom))
-		return false
-	})
-
+// initBalances sets the balance (multiple coins) for an account by address.
+// An error is returned upon failure.
+func (k BaseSendKeeper) initBalances(ctx sdk.Context, addr sdk.AccAddress, balances sdk.Coins) error {
 	accountStore := k.getAccountStore(ctx, addr)
+	for i := range balances {
+		balance := balances[i]
+		if !balance.IsValid() {
+			return sdkerrors.Wrap(sdkerrors.ErrInvalidCoins, balance.String())
+		}
 
-	for _, key := range keys {
-		accountStore.Delete(key)
-	}
-}
-
-// setBalances sets the balance (multiple coins) for an account by address. It will
-// clear out all balances prior to setting the new coins as to set existing balances
-// to zero if they don't exist in amt. An error is returned upon failure.
-func (k BaseSendKeeper) setBalances(ctx sdk.Context, addr sdk.AccAddress, balances sdk.Coins) error {
-	k.clearBalances(ctx, addr)
-
-	for _, balance := range balances {
-		err := k.setBalance(ctx, addr, balance)
-		if err != nil {
-			return err
+		// Bank invariants require to not store zero balances.
+		if !balance.IsZero() {
+			bz := k.cdc.MustMarshal(&balance)
+			accountStore.Set([]byte(balance.Denom), bz)
 		}
 	}