update total liquidity tracker code (#544)

app/app.go

@@ -325,6 +325,12 @@ func NewOsmosisApp(
 
 		})
 
+	app.UpgradeKeeper.SetUpgradeHandler("v5", func(ctx sdk.Context, plan upgradetypes.Plan) {
+		totalLiquidity := app.GAMMKeeper.GetLegacyTotalLiquidity(ctx)
+		app.GAMMKeeper.DeleteLegacyTotalLiquidity(ctx)
+		app.GAMMKeeper.SetTotalLiquidity(ctx, totalLiquidity)
+	})
+
 	// Create IBC Keeper
 	app.IBCKeeper = ibckeeper.NewKeeper(
 		appCodec, keys[ibchost.StoreKey], app.GetSubspace(ibchost.ModuleName), &stakingKeeper, scopedIBCKeeper,

x/gamm/keeper/legacy.go

@@ -0,0 +1,31 @@
+package keeper
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/osmosis-labs/osmosis/x/gamm/types"
+)
+
+func (k Keeper) SetLegacyTotalLiquidity(ctx sdk.Context, coins sdk.Coins) {
+	store := ctx.KVStore(k.storeKey)
+	store.Set(types.KeyTotalLiquidity, []byte(coins.String()))
+}
+
+func (k Keeper) GetLegacyTotalLiquidity(ctx sdk.Context) sdk.Coins {
+	store := ctx.KVStore(k.storeKey)
+	if !store.Has(types.KeyTotalLiquidity) {
+		return sdk.Coins{}
+	}
+
+	bz := store.Get(types.KeyTotalLiquidity)
+	coins, err := sdk.ParseCoinsNormalized(string(bz))
+	if err != nil {
+		panic("invalid total liquidity value set")
+	}
+
+	return coins
+}
+
+func (k Keeper) DeleteLegacyTotalLiquidity(ctx sdk.Context) {
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(types.KeyTotalLiquidity)
+}

x/gamm/keeper/pool_service.go

@@ -3,6 +3,7 @@ package keeper
 import (
 	"fmt"
 
+	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
@@ -559,31 +560,74 @@ func (k Keeper) ExitSwapExternAmountOut(
 }
 
 func (k Keeper) GetTotalLiquidity(ctx sdk.Context) sdk.Coins {
-	store := ctx.KVStore(k.storeKey)
-	if !store.Has(types.KeyTotalLiquidity) {
-		return sdk.Coins{}
+	coins := sdk.Coins{}
+	k.IterateDenomLiquidity(ctx, func(coin sdk.Coin) bool {
+		coins = coins.Add(coin)
+		return false
+	})
+	return coins
+}
+
+func (k Keeper) SetTotalLiquidity(ctx sdk.Context, coins sdk.Coins) {
+	for _, coin := range coins {
+		k.SetDenomLiquidity(ctx, coin.Denom, coin.Amount)
 	}
+}
 
-	bz := store.Get(types.KeyTotalLiquidity)
-	coins, err := sdk.ParseCoinsNormalized(string(bz))
+func (k Keeper) SetDenomLiquidity(ctx sdk.Context, denom string, amount sdk.Int) {
+	store := ctx.KVStore(k.storeKey)
+	bz, err := amount.Marshal()
 	if err != nil {
-		panic("invalid total liquidity value set")
+		panic(err)
 	}
+	store.Set(types.GetDenomPrefix(denom), bz)
+}
 
-	return coins
+func (k Keeper) GetDenomLiquidity(ctx sdk.Context, denom string) sdk.Int {
+	store := ctx.KVStore(k.storeKey)
+	bz := store.Get(types.GetDenomPrefix(denom))
+	if bz == nil {
+		return sdk.NewInt(0)
+	}
+
+	var amount sdk.Int
+	if err := amount.Unmarshal(bz); err != nil {
+		panic(err)
+	}
+	return amount
 }
 
-func (k Keeper) SetTotalLiquidity(ctx sdk.Context, coins sdk.Coins) {
+func (k Keeper) IterateDenomLiquidity(ctx sdk.Context, cb func(sdk.Coin) bool) {
 	store := ctx.KVStore(k.storeKey)
-	store.Set(types.KeyTotalLiquidity, []byte(coins.String()))
+	prefixStore := prefix.NewStore(store, types.KeyTotalLiquidity)
+
+	iterator := prefixStore.Iterator(nil, nil)
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		var amount sdk.Int
+		if err := amount.Unmarshal(iterator.Value()); err != nil {
+			panic(err)
+		}
+
+		if cb(sdk.NewCoin(string(iterator.Key()), amount)) {
+			break
+		}
+	}
 }
 
 func (k Keeper) RecordTotalLiquidityIncrease(ctx sdk.Context, coins sdk.Coins) {
-	liquidity := k.GetTotalLiquidity(ctx)
-	k.SetTotalLiquidity(ctx, liquidity.Add(coins...))
+	for _, coin := range coins {
+		amount := k.GetDenomLiquidity(ctx, coin.Denom)
+		amount = amount.Add(coin.Amount)
+		k.SetDenomLiquidity(ctx, coin.Denom, amount)
+	}
 }
 
 func (k Keeper) RecordTotalLiquidityDecrease(ctx sdk.Context, coins sdk.Coins) {
-	liquidity := k.GetTotalLiquidity(ctx)
-	k.SetTotalLiquidity(ctx, liquidity.Sub(coins))
+	for _, coin := range coins {
+		amount := k.GetDenomLiquidity(ctx, coin.Denom)
+		amount = amount.Sub(coin.Amount)
+		k.SetDenomLiquidity(ctx, coin.Denom, amount)
+	}
 }

x/gamm/types/key.go

@@ -25,6 +25,10 @@ var (
 	KeyTotalLiquidity = []byte{0x03}
 )
 
+func GetDenomPrefix(denom string) []byte {
+	return append(KeyTotalLiquidity, []byte(denom)...)
+}
+
 func GetPoolShareDenom(poolId uint64) string {
 	return fmt.Sprintf("gamm/pool/%d", poolId)
 }