perf: x/bank create reverse prefix for denom<->address

simapp/app_test.go

@@ -123,8 +123,13 @@ func TestRunMigrations(t *testing.T) {
 			false, "", true, "no migrations found for module bank: not found", 0,
 		},
 		{
-			"can register and run migration handler for x/bank",
+			"can register 1->2 migration handler for x/bank, cannot run migration",
 			"bank", 1,
+			false, "", true, "no migration found for module bank from version 2 to version 3: not found", 0,
+		},
+		{
+			"can register 2->3 migration handler for x/bank, can run migration",
+			"bank", 2,
 			false, "", false, "", 1,
 		},
 		{

x/bank/keeper/genesis.go

@@ -13,8 +13,8 @@ func (k BaseKeeper) InitGenesis(ctx sdk.Context, genState *types.GenesisState) {
 	k.SetParams(ctx, genState.Params)
 
 	totalSupply := sdk.Coins{}
-
 	genState.Balances = types.SanitizeGenesisBalances(genState.Balances)
+
 	for _, balance := range genState.Balances {
 		addr, err := sdk.AccAddressFromBech32(balance.Address)
 		if err != nil {

x/bank/keeper/grpc_query.go

@@ -179,24 +179,13 @@ func (k BaseKeeper) DenomOwners(
 	}
 
 	ctx := sdk.UnwrapSDKContext(goCtx)
-
-	store := ctx.KVStore(k.storeKey)
-	balancesStore := prefix.NewStore(store, types.BalancesPrefix)
+	denomPrefixStore := k.getDenomAddressPrefixStore(ctx, req.Denom)
 
 	var denomOwners []*types.DenomOwner
 	pageRes, err := query.FilteredPaginate(
-		balancesStore,
+		denomPrefixStore,
 		req.Pagination,
 		func(key []byte, value []byte, accumulate bool) (bool, error) {
-			var balance sdk.Coin
-			if err := k.cdc.Unmarshal(value, &balance); err != nil {
-				return false, err
-			}
-
-			if req.Denom != balance.Denom {
-				return false, nil
-			}
-
 			if accumulate {
 				address, err := types.AddressFromBalancesStore(key)
 				if err != nil {
@@ -207,7 +196,7 @@ func (k BaseKeeper) DenomOwners(
 					denomOwners,
 					&types.DenomOwner{
 						Address: address.String(),
-						Balance: balance,
+						Balance: k.GetBalance(ctx, address, req.Denom),
 					},
 				)
 			}

x/bank/keeper/migrations.go

@@ -3,6 +3,7 @@ package keeper
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	v043 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v043"
+	v044 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v044"
 )
 
 // Migrator is a struct for handling in-place store migrations.
@@ -19,3 +20,8 @@ func NewMigrator(keeper BaseKeeper) Migrator {
 func (m Migrator) Migrate1to2(ctx sdk.Context) error {
 	return v043.MigrateStore(ctx, m.keeper.storeKey, m.keeper.cdc)
 }
+
+// Migrate2to3 migrates x/bank storage from version 2 to 3.
+func (m Migrator) Migrate2to3(ctx sdk.Context) error {
+	return v044.MigrateStore(ctx, m.keeper.storeKey, m.keeper.cdc)
+}

x/bank/keeper/send.go

@@ -2,8 +2,10 @@ package keeper
 
 import (
 	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
 	"github.com/cosmos/cosmos-sdk/telemetry"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/x/bank/types"
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
@@ -231,16 +233,31 @@ func (k BaseSendKeeper) addCoins(ctx sdk.Context, addr sdk.AccAddress, amt sdk.C
 // 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)
+	denomPrefixStores := make(map[string]prefix.Store) // memoize prefix stores
+
 	for i := range balances {
 		balance := balances[i]
 		if !balance.IsValid() {
 			return sdkerrors.Wrap(sdkerrors.ErrInvalidCoins, balance.String())
 		}
 
-		// Bank invariants require to not store zero balances.
+		// x/bank invariants prohibit persistence of zero balances
 		if !balance.IsZero() {
 			bz := k.cdc.MustMarshal(&balance)
 			accountStore.Set([]byte(balance.Denom), bz)
+
+			denomPrefixStore, ok := denomPrefixStores[balance.Denom]
+			if !ok {
+				denomPrefixStore = k.getDenomAddressPrefixStore(ctx, balance.Denom)
+				denomPrefixStores[balance.Denom] = denomPrefixStore
+			}
+
+			// Store a reverse index from denomination to account address with a
+			// sentinel value.
+			denomAddrKey := address.MustLengthPrefix(addr)
+			if !denomPrefixStore.Has(denomAddrKey) {
+				denomPrefixStore.Set(denomAddrKey, []byte{0})
+			}
 		}
 	}
 
@@ -254,13 +271,22 @@ func (k BaseSendKeeper) setBalance(ctx sdk.Context, addr sdk.AccAddress, balance
 	}
 
 	accountStore := k.getAccountStore(ctx, addr)
+	denomPrefixStore := k.getDenomAddressPrefixStore(ctx, balance.Denom)
 
-	// Bank invariants require to not store zero balances.
+	// x/bank invariants prohibit persistence of zero balances
 	if balance.IsZero() {
 		accountStore.Delete([]byte(balance.Denom))
+		denomPrefixStore.Delete(address.MustLengthPrefix(addr))
 	} else {
 		bz := k.cdc.MustMarshal(&balance)
 		accountStore.Set([]byte(balance.Denom), bz)
+
+		// Store a reverse index from denomination to account address with a
+		// sentinel value.
+		denomAddrKey := address.MustLengthPrefix(addr)
+		if !denomPrefixStore.Has(denomAddrKey) {
+			denomPrefixStore.Set(denomAddrKey, []byte{0})
+		}
 	}
 
 	return nil

x/bank/keeper/view.go

@@ -228,6 +228,11 @@ func (k BaseViewKeeper) ValidateBalance(ctx sdk.Context, addr sdk.AccAddress) er
 // getAccountStore gets the account store of the given address.
 func (k BaseViewKeeper) getAccountStore(ctx sdk.Context, addr sdk.AccAddress) prefix.Store {
 	store := ctx.KVStore(k.storeKey)
-
 	return prefix.NewStore(store, types.CreateAccountBalancesPrefix(addr))
 }
+
+// getDenomAddressPrefixStore returns a prefix store that acts as a reverse index
+// between a denomination and account balance for that denomination.
+func (k BaseViewKeeper) getDenomAddressPrefixStore(ctx sdk.Context, denom string) prefix.Store {
+	return prefix.NewStore(ctx.KVStore(k.storeKey), types.CreateDenomAddressPrefix(denom))
+}

x/bank/migrations/v043/keys.go

@@ -0,0 +1,34 @@
+package v043
+
+import (
+	"errors"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
+)
+
+var (
+	SupplyKey      = []byte{0x00}
+	BalancesPrefix = []byte{0x02}
+
+	ErrInvalidKey = errors.New("invalid key")
+)
+
+func CreateAccountBalancesPrefix(addr []byte) []byte {
+	return append(BalancesPrefix, address.MustLengthPrefix(addr)...)
+}
+
+func AddressFromBalancesStore(key []byte) (sdk.AccAddress, error) {
+	if len(key) == 0 {
+		return nil, ErrInvalidKey
+	}
+
+	addrLen := key[0]
+	bound := int(addrLen)
+
+	if len(key)-1 < bound {
+		return nil, ErrInvalidKey
+	}
+
+	return key[1 : bound+1], nil
+}

x/bank/migrations/v043/store.go

@@ -28,7 +28,7 @@ func migrateSupply(store sdk.KVStore, cdc codec.BinaryCodec) error {
 	}
 
 	// We add a new key for each denom
-	supplyStore := prefix.NewStore(store, types.SupplyKey)
+	supplyStore := prefix.NewStore(store, SupplyKey)
 
 	// We're sure that SupplyI is a Supply struct, there's no other
 	// implementation.
@@ -61,7 +61,7 @@ func migrateBalanceKeys(store sdk.KVStore) {
 	for ; oldStoreIter.Valid(); oldStoreIter.Next() {
 		addr := v040bank.AddressFromBalancesStore(oldStoreIter.Key())
 		denom := oldStoreIter.Key()[v040auth.AddrLen:]
-		newStoreKey := append(types.CreateAccountBalancesPrefix(addr), denom...)
+		newStoreKey := append(CreateAccountBalancesPrefix(addr), denom...)
 
 		// Set new key on store. Values don't change.
 		store.Set(newStoreKey, oldStoreIter.Value())

x/bank/migrations/v044/keys.go

@@ -0,0 +1,10 @@
+package v044
+
+var (
+	DenomAddressPrefix = []byte{0x03}
+)
+
+func CreateAddressDenomPrefix(denom string) []byte {
+	key := append(DenomAddressPrefix, []byte(denom)...)
+	return append(key, 0)
+}

x/bank/migrations/v044/store.go

@@ -0,0 +1,51 @@
+package v044
+
+import (
+	"github.com/cosmos/cosmos-sdk/codec"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
+	v043 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v043"
+)
+
+// MigrateStore performs in-place store migrations from v0.43 to v0.44. The
+// migration includes:
+//
+// - Add an additional reverse index from denomination to address.
+func MigrateStore(ctx sdk.Context, storeKey sdk.StoreKey, cdc codec.BinaryCodec) error {
+	store := ctx.KVStore(storeKey)
+	return addDenomReverseIndex(store, cdc)
+}
+
+func addDenomReverseIndex(store sdk.KVStore, cdc codec.BinaryCodec) error {
+	oldBalancesStore := prefix.NewStore(store, v043.BalancesPrefix)
+
+	oldBalancesIter := oldBalancesStore.Iterator(nil, nil)
+	defer oldBalancesIter.Close()
+
+	denomPrefixStores := make(map[string]prefix.Store) // memoize prefix stores
+
+	for ; oldBalancesIter.Valid(); oldBalancesIter.Next() {
+		var balance sdk.Coin
+		if err := cdc.Unmarshal(oldBalancesIter.Value(), &balance); err != nil {
+			return err
+		}
+
+		addr, err := v043.AddressFromBalancesStore(oldBalancesIter.Key())
+		if err != nil {
+			return err
+		}
+
+		denomPrefixStore, ok := denomPrefixStores[balance.Denom]
+		if !ok {
+			denomPrefixStore = prefix.NewStore(store, CreateAddressDenomPrefix(balance.Denom))
+			denomPrefixStores[balance.Denom] = denomPrefixStore
+		}
+
+		// Store a reverse index from denomination to account address with a
+		// sentinel value.
+		denomPrefixStore.Set(address.MustLengthPrefix(addr), []byte{0})
+	}
+
+	return nil
+}

x/bank/migrations/v044/store_test.go

@@ -0,0 +1,45 @@
+package v044_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"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
+	v043 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v043"
+	v044 "github.com/cosmos/cosmos-sdk/x/bank/migrations/v044"
+)
+
+func TestMigrateStore(t *testing.T) {
+	encCfg := simapp.MakeTestEncodingConfig()
+	bankKey := sdk.NewKVStoreKey("bank")
+	ctx := testutil.DefaultContext(bankKey, sdk.NewTransientStoreKey("transient_test"))
+	store := ctx.KVStore(bankKey)
+
+	addr := sdk.AccAddress([]byte("addr________________"))
+	prefixAccStore := prefix.NewStore(store, v043.CreateAccountBalancesPrefix(addr))
+
+	balances := sdk.NewCoins(
+		sdk.NewCoin("foo", sdk.NewInt(10000)),
+		sdk.NewCoin("bar", sdk.NewInt(20000)),
+	)
+
+	for _, b := range balances {
+		bz, err := encCfg.Codec.Marshal(&b)
+		require.NoError(t, err)
+
+		prefixAccStore.Set([]byte(b.Denom), bz)
+	}
+
+	require.NoError(t, v044.MigrateStore(ctx, bankKey, encCfg.Codec))
+
+	for _, b := range balances {
+		denomPrefixStore := prefix.NewStore(store, v044.CreateAddressDenomPrefix(b.Denom))
+		bz := denomPrefixStore.Get(address.MustLengthPrefix(addr))
+		require.NotNil(t, bz)
+	}
+}

x/bank/module.go

@@ -103,7 +103,13 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 	types.RegisterQueryServer(cfg.QueryServer(), am.keeper)
 
 	m := keeper.NewMigrator(am.keeper.(keeper.BaseKeeper))
-	cfg.RegisterMigration(types.ModuleName, 1, m.Migrate1to2)
+	if err := cfg.RegisterMigration(types.ModuleName, 1, m.Migrate1to2); err != nil {
+		panic(fmt.Sprintf("failed to migrate x/bank from version 1 to 2: %v", err))
+	}
+
+	if err := cfg.RegisterMigration(types.ModuleName, 2, m.Migrate2to3); err != nil {
+		panic(fmt.Sprintf("failed to migrate x/bank from version 2 to 3: %v", err))
+	}
 }
 
 // NewAppModule creates a new AppModule object
@@ -156,7 +162,7 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 }
 
 // ConsensusVersion implements AppModule/ConsensusVersion.
-func (AppModule) ConsensusVersion() uint64 { return 2 }
+func (AppModule) ConsensusVersion() uint64 { return 3 }
 
 // BeginBlock performs a no-op.
 func (AppModule) BeginBlock(_ sdk.Context, _ abci.RequestBeginBlock) {}

x/bank/spec/01_state.md

@@ -4,9 +4,16 @@ order: 1
 
 # State
 
-The `x/bank` module keeps state of three primary objects, account balances, denom metadata and the
-total supply of all balances.
+The `x/bank` module keeps state of three primary objects:
 
-- Supply: `0x0 | byte(denom) -> byte(amount)`
-- Denom Metadata: `0x1 | byte(denom) -> ProtocolBuffer(Metadata)`
-- Balances: `0x2 | byte(address length) | []byte(address) | []byte(balance.Denom) -> ProtocolBuffer(balance)`
+1. Account balances
+2. Denomination metadata
+3. The total supply of all balances
+
+In addition, the `x/bank` module keeps the following indexes to manage the
+aforementioned state:
+
+- Supply Index: `0x0 | byte(denom) -> byte(amount)`
+- Denom Metadata Index: `0x1 | byte(denom) -> ProtocolBuffer(Metadata)`
+- Balances Index: `0x2 | byte(address length) | []byte(address) | []byte(balance.Denom) -> ProtocolBuffer(balance)`
+- Reverse Denomination to Address Index: `0x03 | byte(denom) | 0x00 | []byte(address) -> 0`