feat: add migrations for balances with zero coins (cosmos#9664)

<!--
The default pull request template is for types feat, fix, or refactor.
For other templates, add one of the following parameters to the url:
- template=docs.md
- template=other.md
-->

## Description

Closes: cosmos#9653 

<!-- Add a description of the changes that this PR introduces and the files that
are the most critical to review. -->

---

### Author Checklist

*All items are required. Please add a note to the item if the item is not applicable and
please add links to any relevant follow up issues.*

I have...

- [ ] included the correct [type prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json) in the PR title
- [ ] added `!` to the type prefix if API or client breaking change
- [ ] targeted the correct branch (see [PR Targeting](https://github.com/cosmos/cosmos-sdk/blob/master/CONTRIBUTING.md#pr-targeting))
- [ ] provided a link to the relevant issue or specification
- [ ] followed the guidelines for [building modules](https://github.com/cosmos/cosmos-sdk/blob/master/docs/building-modules)
- [ ] included the necessary unit and integration [tests](https://github.com/cosmos/cosmos-sdk/blob/master/CONTRIBUTING.md#testing)
- [ ] added a changelog entry to `CHANGELOG.md`
- [ ] included comments for [documenting Go code](https://blog.golang.org/godoc)
- [ ] updated the relevant documentation or specification
- [ ] reviewed "Files changed" and left comments if necessary
- [ ] confirmed all CI checks have passed

### Reviewers Checklist

*All items are required. Please add a note if the item is not applicable and please add
your handle next to the items reviewed if you only reviewed selected items.*

I have...

- [ ] confirmed the correct [type prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json) in the PR title
- [ ] confirmed `!` in the type prefix if API or client breaking change
- [ ] confirmed all author checklist items have been addressed 
- [ ] reviewed state machine logic
- [ ] reviewed API design and naming
- [ ] reviewed documentation is accurate
- [ ] reviewed tests and test coverage
- [ ] manually tested (if applicable)

x/bank/migrations/v043/json.go

@@ -0,0 +1,32 @@
+package v043
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/bank/types"
+)
+
+// pruneZeroBalancesJSON removes the zero balance addresses from exported genesis.
+func pruneZeroBalancesJSON(oldBalances []types.Balance) []types.Balance {
+	var balances []types.Balance
+
+	for _, b := range oldBalances {
+		if !b.Coins.IsZero() {
+			b.Coins = sdk.NewCoins(b.Coins...) // prunes zero denom.
+			balances = append(balances, b)
+		}
+	}
+
+	return balances
+}
+
+// MigrateJSON accepts exported v0.40 x/bank genesis state and migrates it to
+// v0.43 x/bank genesis state. The migration includes:
+// - Prune balances & supply with zero coins (ref: https://github.com/cosmos/cosmos-sdk/pull/9229)
+func MigrateJSON(oldState *types.GenesisState) *types.GenesisState {
+	return &types.GenesisState{
+		Params:        oldState.Params,
+		Balances:      pruneZeroBalancesJSON(oldState.Balances),
+		Supply:        sdk.NewCoins(oldState.Supply...), // NewCoins used here to remove zero coin denoms from supply.
+		DenomMetadata: oldState.DenomMetadata,
+	}
+}

x/bank/migrations/v043/json_test.go

@@ -0,0 +1,93 @@
+package v043_test
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/simapp"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	v043bank "github.com/cosmos/cosmos-sdk/x/bank/migrations/v043"
+	"github.com/cosmos/cosmos-sdk/x/bank/types"
+)
+
+func TestMigrateJSON(t *testing.T) {
+	encodingConfig := simapp.MakeTestEncodingConfig()
+	clientCtx := client.Context{}.
+		WithInterfaceRegistry(encodingConfig.InterfaceRegistry).
+		WithTxConfig(encodingConfig.TxConfig).
+		WithCodec(encodingConfig.Codec)
+
+	voter, err := sdk.AccAddressFromBech32("cosmos1fl48vsnmsdzcv85q5d2q4z5ajdha8yu34mf0eh")
+	require.NoError(t, err)
+	bankGenState := &types.GenesisState{
+		Balances: []types.Balance{
+			{
+				Address: voter.String(),
+				Coins: sdk.Coins{
+					sdk.NewCoin("foo", sdk.NewInt(10)),
+					sdk.NewCoin("bar", sdk.NewInt(20)),
+					sdk.NewCoin("foobar", sdk.NewInt(0)),
+				},
+			},
+		},
+		Supply: sdk.Coins{
+			sdk.NewCoin("foo", sdk.NewInt(10)),
+			sdk.NewCoin("bar", sdk.NewInt(20)),
+			sdk.NewCoin("foobar", sdk.NewInt(0)),
+			sdk.NewCoin("barfoo", sdk.NewInt(0)),
+		},
+	}
+
+	migrated := v043bank.MigrateJSON(bankGenState)
+
+	bz, err := clientCtx.Codec.MarshalJSON(migrated)
+	require.NoError(t, err)
+
+	// Indent the JSON bz correctly.
+	var jsonObj map[string]interface{}
+	err = json.Unmarshal(bz, &jsonObj)
+	require.NoError(t, err)
+	indentedBz, err := json.MarshalIndent(jsonObj, "", "\t")
+	require.NoError(t, err)
+
+	// Make sure about:
+	// - zero coin balances pruned.
+	// - zero supply denoms pruned.
+	expected := `{
+	"balances": [
+		{
+			"address": "cosmos1fl48vsnmsdzcv85q5d2q4z5ajdha8yu34mf0eh",
+			"coins": [
+				{
+					"amount": "20",
+					"denom": "bar"
+				},
+				{
+					"amount": "10",
+					"denom": "foo"
+				}
+			]
+		}
+	],
+	"denom_metadata": [],
+	"params": {
+		"default_send_enabled": false,
+		"send_enabled": []
+	},
+	"supply": [
+		{
+			"amount": "20",
+			"denom": "bar"
+		},
+		{
+			"amount": "10",
+			"denom": "foo"
+		}
+	]
+}`
+
+	require.Equal(t, expected, string(indentedBz))
+}

x/bank/migrations/v043/keys.go

@@ -0,0 +1,12 @@
+package v043
+
+const (
+	// ModuleName is the name of the module
+	ModuleName = "bank"
+)
+
+// KVStore keys
+var (
+	BalancesPrefix = []byte{0x02}
+	SupplyKey      = []byte{0x00}
+)

x/bank/migrations/v043/store.go

@@ -75,9 +75,57 @@ func migrateBalanceKeys(store sdk.KVStore) {
 // - Change addresses to be length-prefixed.
 // - Change balances prefix to 1 byte
 // - Change supply to be indexed by denom
+// - Prune balances & supply with zero coins (ref: https://github.com/cosmos/cosmos-sdk/pull/9229)
 func MigrateStore(ctx sdk.Context, storeKey sdk.StoreKey, cdc codec.BinaryCodec) error {
 	store := ctx.KVStore(storeKey)
-
 	migrateBalanceKeys(store)
-	return migrateSupply(store, cdc)
+
+	if err := pruneZeroBalances(store, cdc); err != nil {
+		return err
+	}
+
+	if err := migrateSupply(store, cdc); err != nil {
+		return err
+	}
+
+	return pruneZeroSupply(store)
+}
+
+// pruneZeroBalances removes the zero balance addresses from balances store.
+func pruneZeroBalances(store sdk.KVStore, cdc codec.BinaryCodec) error {
+	balancesStore := prefix.NewStore(store, BalancesPrefix)
+	iterator := balancesStore.Iterator(nil, nil)
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		var balance sdk.Coin
+		if err := cdc.Unmarshal(iterator.Value(), &balance); err != nil {
+			return err
+		}
+
+		if balance.IsZero() {
+			balancesStore.Delete(iterator.Key())
+		}
+	}
+	return nil
+}
+
+// pruneZeroSupply removes zero balance denom from supply store.
+func pruneZeroSupply(store sdk.KVStore) error {
+	supplyStore := prefix.NewStore(store, SupplyKey)
+	iterator := supplyStore.Iterator(nil, nil)
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		var amount sdk.Int
+		if err := amount.Unmarshal(iterator.Value()); err != nil {
+			return err
+		}
+
+		if amount.IsZero() {
+			supplyStore.Delete(iterator.Key())
+		}
+	}
+
+	return nil
 }

x/bank/migrations/v043/store_test.go

@@ -23,10 +23,11 @@ func TestSupplyMigration(t *testing.T) {
 
 	oldFooCoin := sdk.NewCoin("foo", sdk.NewInt(100))
 	oldBarCoin := sdk.NewCoin("bar", sdk.NewInt(200))
+	oldFooBarCoin := sdk.NewCoin("foobar", sdk.NewInt(0)) // to ensure the zero denom coins pruned.
 
 	// Old supply was stored as a single blob under the `SupplyKey`.
 	var oldSupply v040bank.SupplyI
-	oldSupply = &types.Supply{Total: sdk.NewCoins(oldFooCoin, oldBarCoin)}
+	oldSupply = &types.Supply{Total: sdk.Coins{oldFooCoin, oldBarCoin, oldFooBarCoin}}
 	oldSupplyBz, err := encCfg.Codec.MarshalInterface(oldSupply)
 	require.NoError(t, err)
 	store.Set(v040bank.SupplyKey, oldSupplyBz)
@@ -57,6 +58,10 @@ func TestSupplyMigration(t *testing.T) {
 		Amount: amount,
 	}
 	require.Equal(t, oldBarCoin, newBarCoin)
+
+	// foobar shouldn't be existed in the store.
+	bz = supplyStore.Get([]byte("foobar"))
+	require.Nil(t, bz)
 }
 
 func TestBalanceKeysMigration(t *testing.T) {
@@ -66,19 +71,32 @@ func TestBalanceKeysMigration(t *testing.T) {
 	store := ctx.KVStore(bankKey)
 
 	_, _, addr := testdata.KeyTestPubAddr()
-	denom := []byte("foo")
-	value := []byte("bar")
 
-	oldKey := append(append(v040bank.BalancesPrefix, addr...), denom...)
-	store.Set(oldKey, value)
+	// set 10 foo coin
+	fooCoin := sdk.NewCoin("foo", sdk.NewInt(10))
+	oldFooKey := append(append(v040bank.BalancesPrefix, addr...), []byte(fooCoin.Denom)...)
+	fooBz, err := encCfg.Codec.Marshal(&fooCoin)
+	require.NoError(t, err)
+	store.Set(oldFooKey, fooBz)
+
+	// set 0 foobar coin
+	fooBarCoin := sdk.NewCoin("foobar", sdk.NewInt(0))
+	oldKeyFooBar := append(append(v040bank.BalancesPrefix, addr...), []byte(fooBarCoin.Denom)...)
+	fooBarBz, err := encCfg.Codec.Marshal(&fooBarCoin)
+	require.NoError(t, err)
+	store.Set(oldKeyFooBar, fooBarBz)
+	require.NotNil(t, store.Get(oldKeyFooBar)) // before store migation zero values can also exist in store.
 
-	err := v043bank.MigrateStore(ctx, bankKey, encCfg.Codec)
+	err = v043bank.MigrateStore(ctx, bankKey, encCfg.Codec)
 	require.NoError(t, err)
 
-	newKey := append(types.CreateAccountBalancesPrefix(addr), denom...)
+	newKey := append(types.CreateAccountBalancesPrefix(addr), []byte(fooCoin.Denom)...)
 	// -7 because we replaced "balances" with 0x02,
 	// +1 because we added length-prefix to address.
-	require.Equal(t, len(oldKey)-7+1, len(newKey))
-	require.Nil(t, store.Get(oldKey))
-	require.Equal(t, value, store.Get(newKey))
+	require.Equal(t, len(oldFooKey)-7+1, len(newKey))
+	require.Nil(t, store.Get(oldFooKey))
+	require.Equal(t, fooBz, store.Get(newKey))
+
+	newKeyFooBar := append(types.CreateAccountBalancesPrefix(addr), []byte(fooBarCoin.Denom)...)
+	require.Nil(t, store.Get(newKeyFooBar)) // after migration zero balances pruned from store.
 }

x/genutil/migrations/v043/migrate.go

@@ -2,6 +2,9 @@ package v043
 
 import (
 	"github.com/cosmos/cosmos-sdk/client"
+	v040bank "github.com/cosmos/cosmos-sdk/x/bank/migrations/v040"
+	v043bank "github.com/cosmos/cosmos-sdk/x/bank/migrations/v043"
+	bank "github.com/cosmos/cosmos-sdk/x/bank/types"
 	"github.com/cosmos/cosmos-sdk/x/genutil/types"
 	v040gov "github.com/cosmos/cosmos-sdk/x/gov/migrations/v040"
 	v043gov "github.com/cosmos/cosmos-sdk/x/gov/migrations/v043"
@@ -24,5 +27,18 @@ func Migrate(appState types.AppMap, clientCtx client.Context) types.AppMap {
 		appState[v043gov.ModuleName] = clientCtx.Codec.MustMarshalJSON(v043gov.MigrateJSON(&oldGovState))
 	}
 
+	if appState[v040bank.ModuleName] != nil {
+		// unmarshal relative source genesis application state
+		var oldBankState bank.GenesisState
+		clientCtx.Codec.MustUnmarshalJSON(appState[v040bank.ModuleName], &oldBankState)
+
+		// delete deprecated x/bank genesis state
+		delete(appState, v040bank.ModuleName)
+
+		// Migrate relative source genesis application state and marshal it into
+		// the respective key.
+		appState[v043bank.ModuleName] = clientCtx.Codec.MustMarshalJSON(v043bank.MigrateJSON(&oldBankState))
+	}
+
 	return appState
 }