From 9603b15b1a3b444c10e9c59508df84f5a6092085 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20Aasted=20S=C3=B8rensen?= <has@bitcraft.dk>
Date: Thu, 29 Sep 2022 15:41:14 +0200
Subject: [PATCH 1/2] Refactor the calculation of circulating supply to avoid
 iterating all accounts

---
 x/queries/expected_keepers.go |  2 +-
 x/queries/querier.go          | 34 +++++++++++-----------------------
 2 files changed, 12 insertions(+), 24 deletions(-)

diff --git a/x/queries/expected_keepers.go b/x/queries/expected_keepers.go
index baccba2c..642cbc7d 100644
--- a/x/queries/expected_keepers.go
+++ b/x/queries/expected_keepers.go
@@ -13,7 +13,7 @@ type AccountKeeper interface {
 type BankKeeper interface {
 	SpendableCoins(ctx sdk.Context, addr sdk.AccAddress) sdk.Coins
 	GetSupply(ctx sdk.Context) exported.SupplyI
-	IterateAllBalances(ctx sdk.Context, cb func(sdk.AccAddress, sdk.Coin) bool)
+	GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin
 }
 
 type SlashingKeeper interface {
diff --git a/x/queries/querier.go b/x/queries/querier.go
index a819c08f..e14b7bde 100644
--- a/x/queries/querier.go
+++ b/x/queries/querier.go
@@ -42,37 +42,25 @@ func querySpendableBalance(ctx sdk.Context, k BankKeeper, path []string) (res []
 }
 
 func calculateCirculatingSupply(ctx sdk.Context, accK AccountKeeper, bk BankKeeper) (circSupply sdk.Coins) {
-	total := bk.GetSupply(ctx).GetTotal()
-
-	stakingAccounts := map[string]bool{
-		accK.GetModuleAccount(ctx, stakingtypes.NotBondedPoolName).GetAddress().String(): true,
-		accK.GetModuleAccount(ctx, stakingtypes.BondedPoolName).GetAddress().String():    true,
+	stakingAccounts := []sdk.AccAddress{
+		accK.GetModuleAccount(ctx, stakingtypes.NotBondedPoolName).GetAddress(),
+		accK.GetModuleAccount(ctx, stakingtypes.BondedPoolName).GetAddress(),
 	}
 
-	ngmbalance := sdk.ZeroInt()
-
-	bk.IterateAllBalances(ctx, func(address sdk.AccAddress, coin sdk.Coin) bool {
-		if coin.Denom != stakingDenom {
-			return false
-		}
-
-		if _, exists := stakingAccounts[address.String()]; exists {
-			return false
-		}
-
-		spendableCoins := bk.SpendableCoins(ctx, address)
-		ngmbalance = ngmbalance.Add(spendableCoins.AmountOf("ungm"))
-
-		return false
-	})
+	bondedAndUnbondingBalance := sdk.ZeroInt()
+	for _, acc := range stakingAccounts {
+		bal := bk.GetBalance(ctx, acc, "ungm")
+		bondedAndUnbondingBalance = bondedAndUnbondingBalance.Add(bal.Amount)
+	}
 
-	// Replace staking token balance with the one calculated above, which omits vesting and staked tokens.
+	total := bk.GetSupply(ctx).GetTotal()
+	// Replace staking token balance with the one calculated above, which omits staked tokens.
 	for i, c := range total {
 		if c.Denom != stakingDenom {
 			continue
 		}
 
-		total[i] = sdk.NewCoin(stakingDenom, ngmbalance)
+		total[i] = total[i].Sub(sdk.NewCoin("ungm", bondedAndUnbondingBalance))
 		break
 	}
 

From 33cd69be8b6122e0fe5e9aec7d008f379b097fa9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20Aasted=20S=C3=B8rensen?= <has@bitcraft.dk>
Date: Thu, 29 Sep 2022 16:03:05 +0200
Subject: [PATCH 2/2] Fix test of circulating supply query

---
 x/queries/grpc_query_test.go | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/x/queries/grpc_query_test.go b/x/queries/grpc_query_test.go
index 62217a66..a9773a70 100644
--- a/x/queries/grpc_query_test.go
+++ b/x/queries/grpc_query_test.go
@@ -66,7 +66,7 @@ func TestCirculating(t *testing.T) {
 
 	gotRsp, err := queryClient.Circulating(ctx, &types.QueryCirculatingRequest{})
 	require.NoError(t, err)
-	assert.Equal(t, mustParseCoins("154blx,3"+stakingDenom), gotRsp.Total)
+	assert.Equal(t, mustParseCoins("154blx,153"+stakingDenom), gotRsp.Total)
 }
 
 func TestMissedBlocks(t *testing.T) {
@@ -122,6 +122,18 @@ type bankKeeperMock struct {
 	vesting  sdk.Coins
 }
 
+func (b bankKeeperMock) GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin {
+	if bal, found := b.balances[addr.String()]; found {
+		for _, b := range bal {
+			if b.Denom == denom {
+				return b
+			}
+		}
+	}
+
+	return sdk.NewCoin(denom, sdk.ZeroInt())
+}
+
 func (b bankKeeperMock) SpendableCoins(_ sdk.Context, addr sdk.AccAddress) sdk.Coins {
 	return b.balances[addr.String()]
 }
@@ -136,21 +148,6 @@ func (b bankKeeperMock) GetSupply(_ sdk.Context) exported.SupplyI {
 	return banktypes.NewSupply(supply)
 }
 
-func (b bankKeeperMock) IterateAllBalances(_ sdk.Context, cb func(sdk.AccAddress, sdk.Coin) bool) {
-	for address, balance := range b.balances {
-		addr, err := sdk.AccAddressFromBech32(address)
-		if err != nil {
-			panic(err)
-		}
-
-		for _, coin := range balance {
-			if cb(addr, coin) {
-				return
-			}
-		}
-	}
-}
-
 var (
 	_ AccountKeeper  = &accountKeeperMock{}
 	_ BankKeeper     = &bankKeeperMock{}