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

…) (cosmos#9687)

* 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
-->

Closes: cosmos#9653

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

---

*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

*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)

(cherry picked from commit d56c8cd)

* fix conflicts

* fix tests

Co-authored-by: atheeshp <59333759+atheeshp@users.noreply.github.com>
Co-authored-by: atheesh <atheesh@vitwit.com>

x/bank/legacy/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/legacy/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/legacy/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.Marshaler)
+
+	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/legacy/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/legacy/v043/store.go

@@ -0,0 +1,131 @@
+package v043
+
+import (
+	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	v040auth "github.com/cosmos/cosmos-sdk/x/auth/legacy/v040"
+	v040bank "github.com/cosmos/cosmos-sdk/x/bank/legacy/v040"
+	"github.com/cosmos/cosmos-sdk/x/bank/types"
+)
+
+// migrateSupply migrates the supply to be stored by denom key instead in a
+// single blob.
+// ref: https://github.com/cosmos/cosmos-sdk/issues/7092
+func migrateSupply(store sdk.KVStore, cdc codec.BinaryCodec) error {
+	// Old supply was stored as a single blob under the SupplyKey.
+	var oldSupplyI v040bank.SupplyI
+	err := cdc.UnmarshalInterface(store.Get(v040bank.SupplyKey), &oldSupplyI)
+	if err != nil {
+		return err
+	}
+
+	// We delete the single key holding the whole blob.
+	store.Delete(v040bank.SupplyKey)
+
+	if oldSupplyI == nil {
+		return nil
+	}
+
+	// We add a new key for each denom
+	supplyStore := prefix.NewStore(store, types.SupplyKey)
+
+	// We're sure that SupplyI is a Supply struct, there's no other
+	// implementation.
+	oldSupply := oldSupplyI.(*types.Supply)
+	for i := range oldSupply.Total {
+		coin := oldSupply.Total[i]
+		coinBz, err := coin.Amount.Marshal()
+		if err != nil {
+			return err
+		}
+
+		supplyStore.Set([]byte(coin.Denom), coinBz)
+	}
+
+	return nil
+}
+
+// migrateBalanceKeys migrate the balances keys to cater for variable-length
+// addresses.
+func migrateBalanceKeys(store sdk.KVStore) {
+	// old key is of format:
+	// prefix ("balances") || addrBytes (20 bytes) || denomBytes
+	// new key is of format
+	// prefix (0x02) || addrLen (1 byte) || addrBytes || denomBytes
+	oldStore := prefix.NewStore(store, v040bank.BalancesPrefix)
+
+	oldStoreIter := oldStore.Iterator(nil, nil)
+	defer oldStoreIter.Close()
+
+	for ; oldStoreIter.Valid(); oldStoreIter.Next() {
+		addr := v040bank.AddressFromBalancesStore(oldStoreIter.Key())
+		denom := oldStoreIter.Key()[v040auth.AddrLen:]
+		newStoreKey := append(types.CreateAccountBalancesPrefix(addr), denom...)
+
+		// Set new key on store. Values don't change.
+		store.Set(newStoreKey, oldStoreIter.Value())
+		oldStore.Delete(oldStoreIter.Key())
+	}
+}
+
+// MigrateStore performs in-place store migrations from v0.40 to v0.43. The
+// migration includes:
+//
+// - 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)
+
+	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/legacy/v043/store_test.go

@@ -0,0 +1,102 @@
+package v043_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/cosmos/cosmos-sdk/simapp"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	"github.com/cosmos/cosmos-sdk/testutil"
+	"github.com/cosmos/cosmos-sdk/testutil/testdata"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	v040bank "github.com/cosmos/cosmos-sdk/x/bank/legacy/v040"
+	v043bank "github.com/cosmos/cosmos-sdk/x/bank/legacy/v043"
+	"github.com/cosmos/cosmos-sdk/x/bank/types"
+)
+
+func TestSupplyMigration(t *testing.T) {
+	encCfg := simapp.MakeTestEncodingConfig()
+	bankKey := sdk.NewKVStoreKey("bank")
+	ctx := testutil.DefaultContext(bankKey, sdk.NewTransientStoreKey("transient_test"))
+	store := ctx.KVStore(bankKey)
+
+	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.Coins{oldFooCoin, oldBarCoin, oldFooBarCoin}}
+	oldSupplyBz, err := encCfg.Marshaler.MarshalInterface(oldSupply)
+	require.NoError(t, err)
+	store.Set(v040bank.SupplyKey, oldSupplyBz)
+
+	// Run migration.
+	err = v043bank.MigrateStore(ctx, bankKey, encCfg.Marshaler)
+	require.NoError(t, err)
+
+	// New supply is indexed by denom.
+	supplyStore := prefix.NewStore(store, types.SupplyKey)
+	bz := supplyStore.Get([]byte("foo"))
+	var amount sdk.Int
+	err = amount.Unmarshal(bz)
+	require.NoError(t, err)
+
+	newFooCoin := sdk.Coin{
+		Denom:  "foo",
+		Amount: amount,
+	}
+	require.Equal(t, oldFooCoin, newFooCoin)
+
+	bz = supplyStore.Get([]byte("bar"))
+	err = amount.Unmarshal(bz)
+	require.NoError(t, err)
+
+	newBarCoin := sdk.Coin{
+		Denom:  "bar",
+		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) {
+	encCfg := simapp.MakeTestEncodingConfig()
+	bankKey := sdk.NewKVStoreKey("bank")
+	ctx := testutil.DefaultContext(bankKey, sdk.NewTransientStoreKey("transient_test"))
+	store := ctx.KVStore(bankKey)
+
+	_, _, addr := testdata.KeyTestPubAddr()
+
+	// set 10 foo coin
+	fooCoin := sdk.NewCoin("foo", sdk.NewInt(10))
+	oldFooKey := append(append(v040bank.BalancesPrefix, addr...), []byte(fooCoin.Denom)...)
+	fooBz, err := encCfg.Marshaler.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.Marshaler.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.Marshaler)
+	require.NoError(t, err)
+
+	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(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.
+}