[stableswap]: Implement simplified direct multi-asset CFMM solver

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

@@ -152,6 +152,84 @@ func solveCfmmDirect(xReserve, yReserve, yIn osmomath.BigDec) osmomath.BigDec {
 	return xOut
 }
 
+// multi-asset CFMM is xyu(x^2 + y^2 + w) = k
+// As described in our spec, we can ignore the u term and simply solve within the bounds of k' = k / u
+// since u remains constant throughout any independent operation this solver would be used for.
+// We want to solve for a given addition of `b` units of y into the pool,
+// how many units `a` of x do we get out.
+// Let y' = y + b
+// we solve k = (x'y')(x'^2 + y^2 + w) for x', using the following equation: https://www.wolframalpha.com/input?i2d=true&i=solve+for+y%5C%2844%29+x*y*%5C%2840%29Power%5Bx%2C2%5D+%2B+Power%5By%2C2%5D+%2B+w%5C%2841%29%3Dk
+// which we simplify to be the following: https://www.desmos.com/calculator/zx2qslqndl
+// Then we use that to derive the change in x as x_out = x' - x
+//
+// Since original reserves, y' and k are known and remain constant throughout the calculation,
+// deriving x' and then finding x_out is equivalent to finding x_out directly.
+func solveCFMMMultiDirect(xReserve, yReserve, wSumSquares, yIn osmomath.BigDec) osmomath.BigDec {
+	if !xReserve.IsPositive() || !yReserve.IsPositive() || wSumSquares.IsNegative() || !yIn.IsPositive() {
+		panic("invalid input: reserves and input must be positive")
+	} else if yIn.GTE(yReserve) {
+		panic("cannot input more than pool reserves")
+	}
+
+	// find k' using existing reserves (k' = k / v term)
+	k := cfmmConstantMultiNoV(xReserve, yReserve, wSumSquares)
+	k2 := k.Mul(k)
+
+	// find new yReserve after join
+	y_new := yReserve.Add(yIn)
+
+	// store powers to simplify calculations
+	y2 := y_new.Mul(y_new)
+	y3 := y2.Mul(y_new)
+	y4 := y3.Mul(y_new)
+
+	// We then solve for new xReserve using new yReserve and old k using a solver derived from xy(x^2 + y^2 + w) = k
+	// Full equation: x' = (sqrt(729 k^2 y^4 + 108 y^3 (w y + y^3)^3) + 27 k y^2)^(1/3) / (3 2^(1/3) y)
+	// 								- (2^(1/3) (w y + y^3))/(sqrt(729 k^2 y^4 + 108 y^3 (w y + y^3)^3) + 27 k y^2)^(1/3)
+	//
+	//
+	// To simplify, we make the following abstractions:
+	// 1. sqrt_term = sqrt(729 k^2 y^4 + 108 y^3 (w y + y^3)^3)
+	// 2. cube_root_term = (sqrt_term + 27 k y^2)^(1/3)
+	// 3. term1 = cube_root_term / (3 2^(1/3) y)
+	// 4. term2 = (2^(1/3) (w y + y^3)) / cube_root_term
+	//
+	// With these, the final equation becomes: x' = term1 - term2
+
+	// let sqrt_term = sqrt(729 k^2 y^4 + 108 y^3 (w y + y^3)^3)
+	wypy3 := (wSumSquares.Mul(y_new)).Add(y3)
+	wypy3pow3 := wypy3.Mul(wypy3).Mul(wypy3)
+
+	sqrt_term, err := ((k2.Mul(y4).MulInt64(729)).Add(y3.MulInt64(108).Mul(wypy3pow3))).ApproxRoot(2)
+	if err != nil {
+		panic(err)
+	}
+
+	// let cube_root_term = (sqrt_term + 27 k y^2)^(1/3)
+	cube_root_term, err := (sqrt_term.Add(k.Mul(y2).MulInt64(27))).ApproxRoot(3)
+	if err != nil {
+		panic(err)
+	}
+
+	// let term1 = cube_root_term / (3 2^(1/3) y)
+	term1 := cube_root_term.Quo(cubeRootTwo.MulInt64(3).Mul(y_new))
+
+	// let term2 = cube_root_term * (2^(1/3) (w y + y^3))
+	term2 := (cubeRootTwo.Mul(wypy3)).Quo(cube_root_term)
+
+	// finally, let x' = term1 - term2
+	x_new := term1.Sub(term2)
+
+	// find amount of x to output using initial and final xReserve values
+	xOut := xReserve.Sub(x_new)
+
+	if xOut.GTE(xReserve) {
+		panic("invalid output: greater than full pool reserves")
+	}
+
+	return xOut
+}
+
 func approxDecEqual(a, b, tol osmomath.BigDec) bool {
 	return (a.Sub(b).Abs()).LTE(tol)
 }
@@ -197,7 +275,6 @@ 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() {
 		panic("invalid input: reserves and input must be positive")

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

@@ -474,6 +474,30 @@ func TestCFMMInvariantMultiAssets(t *testing.T) {
 	}
 }
 
+func TestCFMMInvariantMultiAssetsDirect(t *testing.T) {
+	kErrTolerance := osmomath.OneDec()
+
+	tests := multiAssetCFMMTestCases
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			// system under test
+			sut := func() {
+				uReserve := calcUReserve(test.remReserves)
+				wSumSquares := calcWSumSquares(test.remReserves)
+
+				// using multi-asset cfmm
+				k2 := cfmmConstantMulti(test.xReserve, test.yReserve, uReserve, wSumSquares)
+				xOut2 := solveCFMMMultiDirect(test.xReserve, test.yReserve, wSumSquares, test.yIn)
+				k3 := cfmmConstantMulti(test.xReserve.Sub(xOut2), test.yReserve.Add(test.yIn), uReserve, wSumSquares)
+				osmomath.DecApproxEq(t, k2, k3, kErrTolerance)
+			}
+
+			osmoassert.ConditionalPanic(t, test.expectPanic, sut)
+		})
+	}
+}
+
 func TestCFMMInvariantMultiAssetsBinarySearch(t *testing.T) {
 	kErrTolerance := osmomath.OneDec()