[x/gamm][stableswap]: Expand inverse relation tests to multi assets and add scaling factors (#3006)

* expand inverse relation tests to multi assets and add scaling factors

* tighten binary search bounds to fit spec

* remove todo

* Fix sped

* move output negation to caller

Co-authored-by: Dev Ojha <dojha@berkeley.edu>

x/gamm/pool-models/stableswap/README.md

@@ -146,7 +146,7 @@ def iterative_search(x_f, y_0, w, k, err_tolerance):
     # k_0 < k. Need to find an upperbound. Worst case assume a linear relationship, gives an upperbound
     # TODO: In the future, we can derive better bounds via reasoning about coefficients in the cubic
     # These are quite close when we are in the "stable" part of the curve though.
-    upperbound = ceil(y_0 * k_ratio)
+    upperbound = ceil(y_0 / k_ratio)
   elif k_ratio > 1:
     # need to find a lowerbound. We could use a cubic relation, but for now we just set it to 0.
     lowerbound = 0

x/gamm/pool-models/stableswap/amm.go

@@ -188,7 +188,7 @@ func solveCFMMBinarySearch(constantFunction func(osmomath.BigDec, osmomath.BigDe
 			panic(err)
 		}
 
-		xOut := xReserve.Sub(x_est).Abs()
+		xOut := xReserve.Sub(x_est)
 		if xOut.GTE(xReserve) {
 			panic("invalid output: greater than full pool reserves")
 		}
@@ -199,15 +199,30 @@ func solveCFMMBinarySearch(constantFunction func(osmomath.BigDec, osmomath.BigDe
 // solveCFMMBinarySearch searches the correct dx using binary search over constant K.
 // added for future extension
 func solveCFMMBinarySearchMulti(xReserve, yReserve, wSumSquares, yIn osmomath.BigDec) osmomath.BigDec {
-	if !xReserve.IsPositive() || !yReserve.IsPositive() || wSumSquares.IsNegative() || !yIn.IsPositive() {
+	if !xReserve.IsPositive() || !yReserve.IsPositive() || wSumSquares.IsNegative() {
 		panic("invalid input: reserves and input must be positive")
-	} else if yIn.GTE(yReserve) {
+	} else if yIn.Abs().GTE(yReserve) {
 		panic("cannot input more than pool reserves")
 	}
-	k := cfmmConstantMultiNoV(xReserve, yReserve, wSumSquares)
 	yFinal := yReserve.Add(yIn)
-	xLowEst := osmomath.ZeroDec()
-	xHighEst := xReserve
+	xLowEst, xHighEst := xReserve, xReserve
+	k0 := cfmmConstantMultiNoV(xReserve, yFinal, wSumSquares)
+	k := cfmmConstantMultiNoV(xReserve, yReserve, wSumSquares)
+	kRatio := k0.Quo(k)
+
+	if kRatio.LT(osmomath.OneDec()) {
+		// k_0 < k. Need to find an upperbound. Worst case assume a linear relationship, gives an upperbound
+		// TODO: In the future, we can derive better bounds via reasoning about coefficients in the cubic
+		// These are quite close when we are in the "stable" part of the curve though.
+		xHighEst = xReserve.Quo(kRatio).Ceil()
+	} else if kRatio.GT(osmomath.OneDec()) {
+		// need to find a lowerbound. We could use a cubic relation, but for now we just set it to 0.
+		xLowEst = osmomath.ZeroDec()
+	} else {
+		// k remains unchanged, so xOut = 0
+		return osmomath.ZeroDec()
+	}
+
 	maxIterations := 256
 	errTolerance := osmoutils.ErrTolerance{AdditiveTolerance: sdk.OneInt(), MultiplicativeTolerance: sdk.Dec{}}
 
@@ -296,6 +311,9 @@ func (p *Pool) calcInAmtGivenOut(tokenOut sdk.Coin, tokenInDenom string, swapFee
 	// We are solving for the amount of token in, cfmm(x,y) = cfmm(x + x_in, y - y_out)
 	// x = tokenInSupply, y = tokenOutSupply, yIn = -tokenOutAmount
 	cfmmIn := solveCfmm(tokenInSupply, tokenOutSupply, remReserves, tokenOutAmount.Neg())
+	// returned cfmmIn is negative, representing we need to add this many tokens to pool.
+	// We invert that negative here.
+	cfmmIn = cfmmIn.Neg()
 	// handle swap fee
 	inAmt := cfmmIn.QuoRoundUp(oneMinus(swapFee))
 	// divide by (1 - swapfee) to force a corresponding increase in input asset

x/gamm/pool-models/stableswap/amm_test.go

@@ -505,41 +505,153 @@ func (suite *StableSwapTestSuite) Test_StableSwap_CalculateAmountOutAndIn_Invers
 
 		denomIn       string
 		initialPoolIn int64
+
+		poolLiquidity  sdk.Coins
+		scalingFactors []uint64
 	}
 
 	// For every test case in testcases, apply a swap fee in swapFeeCases.
-	testcases := []testcase{
-		{
+	testcases := map[string]testcase{
+		// two-asset pools
+		"even pool": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000_000_000)),
+			),
+			scalingFactors: []uint64{1, 1},
+		},
+		"uneven pool (2:1)": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(500_000)),
+			),
+			scalingFactors: []uint64{1, 1},
+		},
+		"uneven pool (1_000_000:1)": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000)),
+			),
+			scalingFactors: []uint64{1, 1},
+		},
+		"uneven pool (1:1_000_000)": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000_000_000)),
+			),
+			scalingFactors: []uint64{1, 1},
+		},
+		"even pool, uneven scaling factors": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000_000_000)),
+			),
+			scalingFactors: []uint64{1, 8},
+		},
+		"uneven pool, uneven scaling factors": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(500_000)),
+			),
+			scalingFactors: []uint64{1, 9},
+		},
+
+		// multi asset pools
+		"even multi-asset pool": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("foo", sdk.NewInt(1_000_000)),
+			),
+			scalingFactors: []uint64{1, 1, 1},
+		},
+		"uneven multi-asset pool (2:1:2)": {
+			denomIn:        "ion",
+			denomOut:       "uosmo",
+			initialCalcOut: 100,
+
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(500_000)),
+				sdk.NewCoin("foo", sdk.NewInt(1_000_000)),
+			),
+			scalingFactors: []uint64{1, 1, 1},
+		},
+		"uneven multi-asset pool (1_000_000:1:1_000_000)": {
+			denomIn:        "ion",
 			denomOut:       "uosmo",
-			initialPoolOut: 1_000_000_000,
 			initialCalcOut: 100,
 
-			denomIn:       "ion",
-			initialPoolIn: 1_000_000_000,
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000)),
+				sdk.NewCoin("foo", sdk.NewInt(1_000_000)),
+			),
+			scalingFactors: []uint64{1, 1, 1},
 		},
-		{
+		"uneven multi-asset pool (1:1_000_000:1_000_000)": {
+			denomIn:        "ion",
 			denomOut:       "uosmo",
-			initialPoolOut: 500_000,
 			initialCalcOut: 100,
 
-			denomIn:       "ion",
-			initialPoolIn: 1_000_000,
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("foo", sdk.NewInt(1_000_000)),
+			),
+			scalingFactors: []uint64{1, 1, 1},
 		},
-		{
+		"even multi-asset pool, uneven scaling factors": {
+			denomIn:        "ion",
 			denomOut:       "uosmo",
-			initialPoolOut: 1_000,
 			initialCalcOut: 100,
 
-			denomIn:       "ion",
-			initialPoolIn: 1_000_000_000,
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("foo", sdk.NewInt(1_000_000)),
+			),
+			scalingFactors: []uint64{5, 3, 9},
 		},
-		{
+		"uneven multi-asset pool (2:1:2), uneven scaling factors": {
+			denomIn:        "ion",
 			denomOut:       "uosmo",
-			initialPoolOut: 1_000_000_000,
 			initialCalcOut: 100,
 
-			denomIn:       "ion",
-			initialPoolIn: 1_000,
+			poolLiquidity: sdk.NewCoins(
+				sdk.NewCoin("ion", sdk.NewInt(1_000_000)),
+				sdk.NewCoin("uosmo", sdk.NewInt(500_000)),
+				sdk.NewCoin("foo", sdk.NewInt(1_000_000)),
+			),
+			scalingFactors: []uint64{100, 76, 33},
 		},
 	}
 
@@ -559,9 +671,8 @@ func (suite *StableSwapTestSuite) Test_StableSwap_CalculateAmountOutAndIn_Invers
 			suite.Run(getTestCaseName(tc, swapFee), func() {
 				ctx := suite.CreateTestContext()
 
-				poolLiquidityIn := sdk.NewInt64Coin(tc.denomOut, tc.initialPoolOut)
-				poolLiquidityOut := sdk.NewInt64Coin(tc.denomIn, tc.initialPoolIn)
-				poolLiquidity := sdk.NewCoins(poolLiquidityIn, poolLiquidityOut)
+				poolLiquidityIn := sdk.NewInt64Coin(tc.denomIn, tc.initialPoolIn)
+				poolLiquidityOut := sdk.NewInt64Coin(tc.denomOut, tc.initialPoolOut)
 
 				swapFeeDec, err := sdk.NewDecFromStr(swapFee)
 				suite.Require().NoError(err)
@@ -570,7 +681,7 @@ func (suite *StableSwapTestSuite) Test_StableSwap_CalculateAmountOutAndIn_Invers
 				suite.Require().NoError(err)
 
 				// TODO: add scaling factors into inverse relationship tests
-				pool := createTestPool(suite.T(), poolLiquidity, swapFeeDec, exitFeeDec, []uint64{1, 1})
+				pool := createTestPool(suite.T(), tc.poolLiquidity, swapFeeDec, exitFeeDec, tc.scalingFactors)
 				suite.Require().NotNil(pool)
 				test_helpers.TestCalculateAmountOutAndIn_InverseRelationship(suite.T(), ctx, pool, poolLiquidityIn.Denom, poolLiquidityOut.Denom, tc.initialCalcOut, swapFeeDec)
 			})