perf: Add a map coins struct to speedup bank genesis (#15764)

Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>
Co-authored-by: Marko <marbar3778@yahoo.com>

CHANGELOG.md

@@ -100,6 +100,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (baseapp) [#15023](https://github.com/cosmos/cosmos-sdk/pull/15023) & [#15213](https://github.com/cosmos/cosmos-sdk/pull/15213) Add `MessageRouter` interface to baseapp and pass it to authz, gov and groups instead of concrete type. 
 * (simtestutil) [#15305](https://github.com/cosmos/cosmos-sdk/pull/15305) Add `AppStateFnWithExtendedCb` with callback function to extend rawState.
 * (x/consensus) [#15553](https://github.com/cosmos/cosmos-sdk/pull/15553) Migrate consensus module to use collections
+* (x/bank) [#15764](https://github.com/cosmos/cosmos-sdk/pull/15764) Speedup x/bank InitGenesis
 * (x/auth) [#15867](https://github.com/cosmos/cosmos-sdk/pull/15867) Support better logging for signature verification failure.
 
 ### State Machine Breaking

types/coin_benchmark_test.go

@@ -70,3 +70,71 @@ func BenchmarkCoinsAdditionNoIntersect(b *testing.B) {
 		b.Run(fmt.Sprintf("sizes: A_%d, B_%d", sizeA, sizeB), benchmarkingFunc(sizeA, sizeB))
 	}
 }
+
+func BenchmarkSumOfCoinAdds(b *testing.B) {
+	// This benchmark tests the performance of adding a large number of coins
+	// into a single coin set.
+	// it does numAdds additions, each addition has (numIntersectingCoins) that contain denoms
+	// already in the sum, and (coinsPerAdd - numIntersectingCoins) that are new denoms.
+	benchmarkingFunc := func(numAdds, coinsPerAdd, numIntersectingCoins int, sumFn func([]Coins) Coins) func(b *testing.B) {
+		return func(b *testing.B) {
+			b.ReportAllocs()
+			addCoins := make([]Coins, numAdds)
+			nonIntersectingCoins := coinsPerAdd - numIntersectingCoins
+
+			for i := 0; i < numAdds; i++ {
+				intersectCoins := make([]Coin, numIntersectingCoins)
+				num := NewInt(int64(i))
+				for j := 0; j < numIntersectingCoins; j++ {
+					intersectCoins[j] = NewCoin(coinName(j+1_000_000_000), num)
+				}
+				addCoins[i] = intersectCoins
+				for j := 0; j < nonIntersectingCoins; j++ {
+					addCoins[i] = addCoins[i].Add(NewCoin(coinName(i*nonIntersectingCoins+j), num))
+				}
+			}
+
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				sumFn(addCoins)
+			}
+		}
+	}
+
+	MapCoinsSumFn := func(coins []Coins) Coins {
+		sum := MapCoins{}
+		for _, coin := range coins {
+			sum.Add(coin...)
+		}
+		return sum.ToCoins()
+	}
+	CoinsSumFn := func(coins []Coins) Coins {
+		sum := Coins{}
+		for _, coin := range coins {
+			sum = sum.Add(coin...)
+		}
+		return sum
+	}
+
+	// larger benchmarks with non-overlapping coins won't terminate in reasonable timeframes with sdk.Coins
+	// they work fine with MapCoins
+	benchmarkSizes := [][]int{{5, 2, 1000}, {10, 10, 10000}}
+	sumFns := []struct {
+		name string
+		fn   func([]Coins) Coins
+	}{
+		{"MapCoins", MapCoinsSumFn}, {"Coins", CoinsSumFn},
+	}
+	for i := 0; i < len(benchmarkSizes); i++ {
+		for j := 0; j < 2; j++ {
+			coinsPerAdd := benchmarkSizes[i][0]
+			intersectingCoinsPerAdd := benchmarkSizes[i][1]
+			numAdds := benchmarkSizes[i][2]
+			sumFn := sumFns[j]
+			b.Run(fmt.Sprintf("Fn: %s, num adds: %d, coinsPerAdd: %d, intersecting: %d",
+				sumFn.name, numAdds, coinsPerAdd, intersectingCoinsPerAdd),
+				benchmarkingFunc(numAdds, coinsPerAdd, intersectingCoinsPerAdd, sumFn.fn))
+		}
+	}
+}

types/mapcoins.go

@@ -0,0 +1,42 @@
+package types
+
+// map coins is a map representation of sdk.Coins
+// intended solely for use in bulk additions.
+// All serialization and iteration should be done after conversion to sdk.Coins.
+type MapCoins map[string]Int
+
+func NewMapCoins(coins Coins) MapCoins {
+	m := make(MapCoins, len(coins))
+	m.Add(coins...)
+	return m
+}
+
+func (m MapCoins) Add(coins ...Coin) {
+	for _, coin := range coins {
+		existAmt, exists := m[coin.Denom]
+		// TODO: Once int supports in-place arithmetic, switch this to be in-place.
+		if exists {
+			m[coin.Denom] = existAmt.Add(coin.Amount)
+		} else {
+			m[coin.Denom] = coin.Amount
+		}
+	}
+}
+
+func (m MapCoins) ToCoins() Coins {
+	if len(m) == 0 {
+		return Coins{}
+	}
+	coins := make(Coins, 0, len(m))
+	for denom, amount := range m {
+		if amount.IsZero() {
+			continue
+		}
+		coins = append(coins, NewCoin(denom, amount))
+	}
+	if len(coins) == 0 {
+		return Coins{}
+	}
+	coins.Sort()
+	return coins
+}

types/mapcoins_test.go

@@ -0,0 +1,62 @@
+package types_test
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/stretchr/testify/require"
+)
+
+func (s *coinTestSuite) TestMapCoinsAdd() {
+	cA0M0 := sdk.Coins{s.ca0, s.cm0}
+	cA0M1 := sdk.Coins{s.ca0, s.cm1}
+	cA1M1 := sdk.Coins{s.ca1, s.cm1}
+	cases := []struct {
+		name     string
+		inputOne sdk.Coins
+		inputTwo sdk.Coins
+		expected sdk.Coins
+		msg      string
+	}{
+		{"adding two empty lists", s.emptyCoins, s.emptyCoins, s.emptyCoins, "empty, non list should be returned"},
+		{"empty list + set", s.emptyCoins, cA0M1, sdk.Coins{s.cm1}, "zero coins should be removed"},
+		{"empty list + set", s.emptyCoins, cA1M1, cA1M1, "zero + a_non_zero = a_non_zero"},
+		{"set + empty list", cA0M1, s.emptyCoins, sdk.Coins{s.cm1}, "zero coins should be removed"},
+		{"set + empty list", cA1M1, s.emptyCoins, cA1M1, "a_non_zero + zero  = a_non_zero"},
+		{
+			"{1atom,1muon}+{1atom,1muon}", cA1M1, cA1M1,
+			sdk.Coins{s.ca2, s.cm2},
+			"a + a = 2a",
+		},
+		{
+			"{0atom,1muon}+{0atom,0muon}", cA0M1, cA0M0,
+			sdk.Coins{s.cm1},
+			"zero coins should be removed",
+		},
+		{
+			"{2atom}+{0muon}",
+			sdk.Coins{s.ca2},
+			sdk.Coins{s.cm0},
+			sdk.Coins{s.ca2},
+			"zero coins should be removed",
+		},
+		{
+			"{1atom}+{1atom,2muon}",
+			sdk.Coins{s.ca1},
+			sdk.Coins{s.ca1, s.cm2},
+			sdk.Coins{s.ca2, s.cm2},
+			"should be correctly added",
+		},
+		{
+			"{0atom,0muon}+{0atom,0muon}", cA0M0, cA0M0, s.emptyCoins,
+			"sets with zero coins should return empty set",
+		},
+	}
+
+	for _, tc := range cases {
+		expected := tc.inputOne.Add(tc.inputTwo...)
+		m := sdk.NewMapCoins(tc.inputOne)
+		m.Add(tc.inputTwo...)
+		res := m.ToCoins()
+		s.Require().True(res.IsValid())
+		require.Equal(s.T(), expected, res, tc.msg)
+	}
+}

x/bank/keeper/genesis.go

@@ -19,8 +19,8 @@ func (k BaseKeeper) InitGenesis(ctx sdk.Context, genState *types.GenesisState) {
 	for _, se := range genState.GetAllSendEnabled() {
 		k.SetSendEnabled(ctx, se.Denom, se.Enabled)
 	}
+	totalSupplyMap := sdk.NewMapCoins(sdk.Coins{})
 
-	totalSupply := sdk.Coins{}
 	genState.Balances = types.SanitizeGenesisBalances(genState.Balances)
 
 	for _, balance := range genState.Balances {
@@ -33,8 +33,9 @@ func (k BaseKeeper) InitGenesis(ctx sdk.Context, genState *types.GenesisState) {
 			}
 		}
 
-		totalSupply = totalSupply.Add(balance.Coins...)
+		totalSupplyMap.Add(balance.Coins...)
 	}
+	totalSupply := totalSupplyMap.ToCoins()
 
 	if !genState.Supply.Empty() && !genState.Supply.Equal(totalSupply) {
 		panic(fmt.Errorf("genesis supply is incorrect, expected %v, got %v", genState.Supply, totalSupply))