Improve set supply

x/bank/keeper/genesis.go

@@ -31,7 +31,9 @@ func (k BaseKeeper) InitGenesis(ctx sdk.Context, genState *types.GenesisState) {
 		panic(fmt.Errorf("genesis supply is incorrect, expected %v, got %v", genState.Supply, totalSupply))
 	}
 
-	k.setSupply(ctx, totalSupply)
+	for _, supply := range totalSupply {
+		k.setSupply(ctx, supply)
+	}
 
 	for _, meta := range genState.DenomMetadata {
 		k.SetDenomMetaData(ctx, meta)

x/bank/keeper/keeper.go

@@ -184,33 +184,6 @@ func (k BaseKeeper) GetSupply(ctx sdk.Context, denom string) sdk.Coin {
 	return coin
 }
 
-// SetSupply sets the Supply to store
-func (k BaseKeeper) setSupply(ctx sdk.Context, supply sdk.Coins) {
-	store := ctx.KVStore(k.storeKey)
-	supplyStore := prefix.NewStore(store, types.SupplyKey)
-
-	var newSupply []sdk.Coin
-	storeSupply := k.GetTotalSupply(ctx)
-
-	// update supply for coins which have non zero amount
-	for _, coin := range storeSupply {
-		if supply.AmountOf(coin.Denom).IsZero() {
-			zeroCoin := &sdk.Coin{
-				Denom:  coin.Denom,
-				Amount: sdk.NewInt(0),
-			}
-			bz := k.cdc.MustMarshalBinaryBare(zeroCoin)
-			supplyStore.Set([]byte(coin.Denom), bz)
-		}
-	}
-	newSupply = append(newSupply, supply...)
-
-	for i := range newSupply {
-		bz := k.cdc.MustMarshalBinaryBare(&supply[i])
-		supplyStore.Set([]byte(supply[i].Denom), bz)
-	}
-}
-
 // GetDenomMetaData retrieves the denomination metadata
 func (k BaseKeeper) GetDenomMetaData(ctx sdk.Context, denom string) (types.Metadata, bool) {
 	store := ctx.KVStore(k.storeKey)
@@ -354,7 +327,7 @@ func (k BaseKeeper) UndelegateCoinsFromModuleToAccount(
 
 // MintCoins creates new coins from thin air and adds it to the module account.
 // It will panic if the module account does not exist or is unauthorized.
-func (k BaseKeeper) MintCoins(ctx sdk.Context, moduleName string, amt sdk.Coins) error {
+func (k BaseKeeper) MintCoins(ctx sdk.Context, moduleName string, amounts sdk.Coins) error {
 	acc := k.ak.GetModuleAccount(ctx, moduleName)
 	if acc == nil {
 		panic(sdkerrors.Wrapf(sdkerrors.ErrUnknownAddress, "module account %s does not exist", moduleName))
@@ -364,31 +337,32 @@ func (k BaseKeeper) MintCoins(ctx sdk.Context, moduleName string, amt sdk.Coins)
 		panic(sdkerrors.Wrapf(sdkerrors.ErrUnauthorized, "module account %s does not have permissions to mint tokens", moduleName))
 	}
 
-	err := k.addCoins(ctx, acc.GetAddress(), amt)
+	err := k.addCoins(ctx, acc.GetAddress(), amounts)
 	if err != nil {
 		return err
 	}
 
 	// update total supply
-	supply := k.GetTotalSupply(ctx)
-	supply = supply.Add(amt...)
-
-	k.setSupply(ctx, supply)
+	for _, amount := range amounts {
+		supply := k.GetSupply(ctx, amount.GetDenom())
+		supply = supply.Add(amount)
+		k.setSupply(ctx, supply)
+	}
 
 	logger := k.Logger(ctx)
-	logger.Info("minted coins from module account", "amount", amt.String(), "from", moduleName)
+	logger.Info("minted coins from module account", "amount", amounts.String(), "from", moduleName)
 
 	// emit mint event
 	ctx.EventManager().EmitEvent(
-		types.NewCoinMintEvent(acc.GetAddress(), amt),
+		types.NewCoinMintEvent(acc.GetAddress(), amounts),
 	)
 
 	return nil
 }
 
 // BurnCoins burns coins deletes coins from the balance of the module account.
 // It will panic if the module account does not exist or is unauthorized.
-func (k BaseKeeper) BurnCoins(ctx sdk.Context, moduleName string, amt sdk.Coins) error {
+func (k BaseKeeper) BurnCoins(ctx sdk.Context, moduleName string, amounts sdk.Coins) error {
 	acc := k.ak.GetModuleAccount(ctx, moduleName)
 	if acc == nil {
 		panic(sdkerrors.Wrapf(sdkerrors.ErrUnknownAddress, "module account %s does not exist", moduleName))
@@ -398,28 +372,36 @@ func (k BaseKeeper) BurnCoins(ctx sdk.Context, moduleName string, amt sdk.Coins)
 		panic(sdkerrors.Wrapf(sdkerrors.ErrUnauthorized, "module account %s does not have permissions to burn tokens", moduleName))
 	}
 
-	err := k.subUnlockedCoins(ctx, acc.GetAddress(), amt)
+	err := k.subUnlockedCoins(ctx, acc.GetAddress(), amounts)
 	if err != nil {
 		return err
 	}
 
-	// update total supply
-	supply := k.GetTotalSupply(ctx)
-	supply = supply.Sub(amt)
-
-	k.setSupply(ctx, supply)
+	for _, amount := range amounts {
+		supply := k.GetSupply(ctx, amount.GetDenom())
+		supply = supply.Sub(amount)
+		k.setSupply(ctx, supply)
+	}
 
 	logger := k.Logger(ctx)
-	logger.Info("burned tokens from module account", "amount", amt.String(), "from", moduleName)
+	logger.Info("burned tokens from module account", "amount", amounts.String(), "from", moduleName)
 
 	// emit burn event
 	ctx.EventManager().EmitEvent(
-		types.NewCoinBurnEvent(acc.GetAddress(), amt),
+		types.NewCoinBurnEvent(acc.GetAddress(), amounts),
 	)
 
 	return nil
 }
 
+func (k BaseKeeper) setSupply(ctx sdk.Context, coin sdk.Coin) {
+	store := ctx.KVStore(k.storeKey)
+	supplyStore := prefix.NewStore(store, types.SupplyKey)
+
+	bz := k.cdc.MustMarshalBinaryBare(&coin)
+	supplyStore.Set([]byte(coin.GetDenom()), bz)
+}
+
 func (k BaseKeeper) trackDelegation(ctx sdk.Context, addr sdk.AccAddress, balance, amt sdk.Coins) error {
 	acc := k.ak.GetAccount(ctx, addr)
 	if acc == nil {