diff --git a/app/app.go b/app/app.go index 53b95c2dd93..50732fdb6de 100644 --- a/app/app.go +++ b/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, diff --git a/x/gamm/keeper/legacy.go b/x/gamm/keeper/legacy.go new file mode 100644 index 00000000000..b78241ffe67 --- /dev/null +++ b/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) +} diff --git a/x/gamm/keeper/pool_service.go b/x/gamm/keeper/pool_service.go index 241ee293198..fd39e7ff8db 100644 --- a/x/gamm/keeper/pool_service.go +++ b/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) + } } diff --git a/x/gamm/types/key.go b/x/gamm/types/key.go index 4130b153008..a148882d269 100644 --- a/x/gamm/types/key.go +++ b/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) }