feat: more error handling for updateSwapInvariant (#1764)

* feat: add error handling for updateSwapInvariant reducing base below total short

* chore: changelog

* fix: fix tests

---------

Co-authored-by: Unique Divine <[email protected]>

CHANGELOG.md

@@ -86,6 +86,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#1735](https://github.com/NibiruChain/nibiru/pull/1735) - test(sim): fix simulation tests
 * [#1754](https://github.com/NibiruChain/nibiru/pull/1754) - refactor(decode-base64): clean code improvements and fn docs
 * [#1736](https://github.com/NibiruChain/nibiru/pull/1736) - test(sim): add sim genesis state for all cusom modules
+* [#1764](https://github.com/NibiruChain/nibiru/pull/1764) - fix(perp): make updateswapinvariant aware of total short supply to avoid panics
 
 ### Dependencies
 - Bump `github.com/cometbft/cometbft-db` from 0.9.0 to 0.9.1 ([#1760](https://github.com/NibiruChain/nibiru/pull/1760))

x/perp/v2/integration/action/market.go

@@ -1,6 +1,8 @@
 package action
 
 import (
+	"fmt"
+
 	sdkmath "cosmossdk.io/math"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
@@ -159,6 +161,28 @@ func ShiftSwapInvariant(pair asset.Pair, newValue sdkmath.Int) action.Action {
 	}
 }
 
+type shiftSwapInvariantFail struct {
+	pair     asset.Pair
+	newValue sdkmath.Int
+}
+
+func (e shiftSwapInvariantFail) Do(app *app.NibiruApp, ctx sdk.Context) (sdk.Context, error) {
+	err := app.PerpKeeperV2.Sudo().ShiftSwapInvariant(
+		ctx, e.pair, e.newValue, testapp.DefaultSudoRoot(),
+	)
+	if err == nil {
+		return ctx, fmt.Errorf("expected error")
+	}
+	return ctx, nil
+}
+
+func ShiftSwapInvariantFail(pair asset.Pair, newValue sdkmath.Int) action.Action {
+	return shiftSwapInvariantFail{
+		pair:     pair,
+		newValue: newValue,
+	}
+}
+
 type createPool struct {
 	pair   asset.Pair
 	market types.Market

x/perp/v2/keeper/clearing_house_test.go

@@ -51,14 +51,7 @@ func TestMarketOrder(t *testing.T) {
 				MarketOrder(alice, pairBtcNusd, types.Direction_SHORT, sdk.NewInt(10_000), sdk.OneDec(), sdk.ZeroDec()),
 				MarketOrder(bob, pairBtcNusd, types.Direction_LONG, sdk.NewInt(10_000), sdk.OneDec(), sdk.ZeroDec()),
 
-				ShiftSwapInvariant(pairBtcNusd, sdk.NewInt(1)),
-			).
-			When(
-				PartialCloseFails(alice, pairBtcNusd, sdk.NewDec(5_000), types.ErrAmmNonpositiveReserves),
-			).
-			Then(
-				ClosePosition(bob, pairBtcNusd),
-				PartialClose(alice, pairBtcNusd, sdk.NewDec(5_000)),
+				ShiftSwapInvariantFail(pairBtcNusd, sdk.NewInt(1)),
 			),
 
 		TC("new long position").

x/perp/v2/types/amm.go

@@ -3,6 +3,7 @@ package types
 import (
 	"math/big"
 
+	sdkerrors "cosmossdk.io/errors"
 	sdkmath "cosmossdk.io/math"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -38,6 +39,13 @@ func (amm AMM) Validate() error {
 			"computed sqrt and current sqrt are mismatched. pool: " + amm.String())
 	}
 
+	// Short positions borrow base asset, so if the base reserve is below the
+	// quote reserve swapped to base, then the shorts can't close positions.
+	_, err = amm.SwapBaseAsset(amm.TotalShort, Direction_LONG)
+	if err != nil {
+		return sdkerrors.Wrapf(ErrAmmBaseBorrowedTooHigh, "Base amount error, short exceed total base supply: %s", err.Error())
+	}
+
 	return nil
 }
 
@@ -73,7 +81,7 @@ func (amm AMM) ComputeSettlementPrice() (sdk.Dec, AMM, error) {
 	return price, amm, err
 }
 
-// QuoteReserveToAsset converts quote reserves to assets
+// QuoteReserveToAsset converts quote reserves to assets\
 func (amm AMM) QuoteReserveToAsset(quoteReserve sdk.Dec) sdk.Dec {
 	return QuoteReserveToAsset(quoteReserve, amm.PriceMultiplier)
 }
@@ -413,14 +421,32 @@ func (amm *AMM) UpdateSwapInvariant(newSwapInvariant sdk.Dec) (err error) {
 	// k = x * y
 	// newK = (cx) * (cy) = c^2 xy = c^2 k
 	// newPrice = (c y) / (c x) = y / x = price | unchanged price
-	newSqrtDepth := common.MustSqrtDec(newSwapInvariant)
+	newSqrtDepth, err := common.SqrtDec(newSwapInvariant)
+	if err != nil {
+		return err
+	}
+
 	multiplier := newSqrtDepth.Quo(amm.SqrtDepth)
+	updatedBaseReserve := amm.BaseReserve.Mul(multiplier)
+	updatedQuoteReserve := amm.QuoteReserve.Mul(multiplier)
+
+	newAmm := AMM{
+		BaseReserve:     updatedBaseReserve,
+		QuoteReserve:    updatedQuoteReserve,
+		PriceMultiplier: amm.PriceMultiplier,
+		SqrtDepth:       newSqrtDepth,
+		TotalLong:       amm.TotalLong,
+		TotalShort:      amm.TotalShort,
+	}
+	if err = newAmm.Validate(); err != nil {
+		return err
+	}
 
 	// Change the swap invariant while holding price constant.
 	// Multiplying by the same factor to both of the reserves won't affect price.
 	amm.SqrtDepth = newSqrtDepth
-	amm.BaseReserve = amm.BaseReserve.Mul(multiplier)
-	amm.QuoteReserve = amm.QuoteReserve.Mul(multiplier)
+	amm.BaseReserve = updatedBaseReserve
+	amm.QuoteReserve = updatedQuoteReserve
 
-	return amm.Validate() // might be useless
+	return nil
 }

x/perp/v2/types/amm_test.go

@@ -367,6 +367,8 @@ func TestUpdateSwapInvariant(t *testing.T) {
 				QuoteReserve:    sdk.NewDec(1e6),
 				SqrtDepth:       sdk.NewDec(1e6),
 				PriceMultiplier: sdk.OneDec(),
+				TotalLong:       sdk.ZeroDec(),
+				TotalShort:      sdk.ZeroDec(),
 			},
 			newSwapInvariant:     sdk.NewDec(1e12),
 			expectedBaseReserve:  sdk.NewDec(1e6),
@@ -380,6 +382,8 @@ func TestUpdateSwapInvariant(t *testing.T) {
 				QuoteReserve:    sdk.NewDec(1e6),
 				SqrtDepth:       sdk.NewDec(1e6),
 				PriceMultiplier: sdk.OneDec(),
+				TotalLong:       sdk.ZeroDec(),
+				TotalShort:      sdk.ZeroDec(),
 			},
 			newSwapInvariant:     sdk.NewDec(1e14),
 			expectedBaseReserve:  sdk.NewDec(1e7),
@@ -393,6 +397,8 @@ func TestUpdateSwapInvariant(t *testing.T) {
 				QuoteReserve:    sdk.NewDec(1e6),
 				SqrtDepth:       sdk.NewDec(1e6),
 				PriceMultiplier: sdk.OneDec(),
+				TotalLong:       sdk.ZeroDec(),
+				TotalShort:      sdk.ZeroDec(),
 			},
 			newSwapInvariant:     sdk.NewDec(1e10),
 			expectedBaseReserve:  sdk.NewDec(1e5),
@@ -768,6 +774,34 @@ func TestValidateAMM(t *testing.T) {
 	}
 }
 
+func TestValidateAMMShortVsQuote(t *testing.T) {
+	amm := types.AMM{
+		BaseReserve:     sdk.OneDec(),
+		QuoteReserve:    sdk.OneDec(),
+		PriceMultiplier: sdk.OneDec(),
+		SqrtDepth:       sdk.OneDec(),
+		TotalLong:       sdk.ZeroDec(),
+		TotalShort:      sdk.ZeroDec(),
+	}
+
+	err := amm.Validate()
+	require.NoError(t, err)
+
+	quoteAssetDelta, err := amm.SwapBaseAsset(sdk.NewDec(2), types.Direction_SHORT)
+	require.NoError(t, err)
+	assert.Equal(t, sdk.MustNewDecFromStr("0.666666666666666667"), quoteAssetDelta)
+
+	previousAmm := amm
+
+	err = amm.UpdateSwapInvariant(sdk.MustNewDecFromStr("0.01"))
+	require.ErrorContains(t, err, "Base amount error, short exceed total base supply")
+
+	err = amm.UpdateSwapInvariant(sdk.MustNewDecFromStr("-1"))
+	require.ErrorContains(t, err, "square root of negative number")
+
+	require.Equal(t, amm, previousAmm)
+}
+
 func TestPositionNotionalFail(t *testing.T) {
 	amm := mock.TestAMM(sdk.NewDec(-1), sdk.NewDec(2))
 	_, err := amm.GetQuoteReserveAmt(sdk.NewDec(-1), types.Direction_LONG)

x/perp/v2/types/errors.go

@@ -25,6 +25,7 @@ var (
 	ErrAmmNonpositiveReserves      = errorAmm("base and quote reserves must always be positive")
 	ErrLiquidityDepth              = errorAmm("liquidity depth must be positive and equal to the square of the reserves")
 	ErrAmmBaseSupplyNonpositive    = errorAmm("base supply must be > 0")
+	ErrAmmBaseBorrowedTooHigh      = errorAmm("short supply exceeds existing supply")
 	ErrAmmQuoteSupplyNonpositive   = errorAmm("quote supply must be > 0")
 	ErrAmmLiquidityDepthOverflow   = errorAmm("liquidty depth overflow")
 	ErrAmmNonPositivePegMult       = errorAmm("peg multiplier must be > 0")

x/perp/v2/types/genesis_test.go

@@ -36,6 +36,8 @@ func TestGenesisValidate(t *testing.T) {
 		QuoteReserve:    sdk.OneDec(),
 		PriceMultiplier: sdk.OneDec(),
 		SqrtDepth:       sdk.OneDec(),
+		TotalLong:       sdk.ZeroDec(),
+		TotalShort:      sdk.ZeroDec(),
 	}
 	validPositions := types.GenesisPosition{
 		Pair:    pair,
@@ -70,6 +72,8 @@ func TestGenesisValidate(t *testing.T) {
 		QuoteReserve:    sdk.OneDec(),
 		PriceMultiplier: sdk.OneDec(),
 		SqrtDepth:       sdk.OneDec(),
+		TotalLong:       sdk.ZeroDec(),
+		TotalShort:      sdk.ZeroDec(),
 	}
 	invalidPositions := types.GenesisPosition{
 		Pair:    pair,