Add boilerplate logic for stableswap swaps and spot price (#1287)

Adds logic for stableswap implementing the swap and spot price interfaces. The PR is marked as draft until some test cases are added.
______

For contributor use:o

- [ ] Targeted PR against correct branch (see [CONTRIBUTING.md](https://github.com/cosmos/gaia/blob/main/CONTRIBUTING.md#pr-targeting))
- [ ] Updated relevant documentation (`docs/`) or specification (`x/<module>/spec/`)
- [ ] Added a relevant changelog entry to the `Unreleased` section in `CHANGELOG.md`
- [ ] Re-reviewed `Files changed` in the Github PR explorer

x/gamm/keeper/multihop_test.go

@@ -164,7 +164,7 @@ func (suite *KeeperTestSuite) TestBalancerPoolSimpleMultihopSwapExactAmountOut()
 			}()
 
 			tokenInAmount, err := keeper.MultihopSwapExactAmountOut(suite.Ctx, suite.TestAccs[0], test.param.routes, test.param.tokenInMaxAmount, test.param.tokenOut)
-			suite.NoError(err, "test: %v", test.name)
+			suite.Require().NoError(err, "test: %v", test.name)
 
 			// Calculate the chained spot price.
 			spotPriceAfter := func() sdk.Dec {

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

@@ -119,6 +119,11 @@ func (p Pool) CalcInAmtGivenOut(
 	// Thus in order to give X amount out, we solve the invariant for the invariant input. However invariant input = (1 - swapfee) * trade input.
 	// Therefore we divide by (1 - swapfee) here
 	tokenAmountInBeforeFee := tokenAmountIn.Quo(sdk.OneDec().Sub(swapFee))
+	// TODO: Once we make Calc methods return integers
+	// if tokenInDecimal is non-zero, we add 1 to the tokenInCoin
+	// if tokenInDecimal.Amount.IsPositive() {
+	// 	tokenInCoin.Amount = tokenInCoin.Amount.AddRaw(1)
+	// }
 	return sdk.NewDecCoinFromDec(tokenInDenom, tokenAmountInBeforeFee), nil
 }
 
@@ -132,6 +137,7 @@ func (p *Pool) SwapInAmtGivenOut(
 		return sdk.Coin{}, sdkerrors.Wrapf(types.ErrInvalidMathApprox, "token amount is zero or negative")
 	}
 	tokenInCoin, _ := tokenInDecCoin.TruncateDecimal()
+
 	if !tokenInCoin.Amount.IsPositive() {
 		return sdk.Coin{}, sdkerrors.Wrapf(types.ErrInvalidMathApprox, "token amount must be positive")
 	}

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

@@ -188,3 +188,31 @@ func spotPrice(baseReserve, quoteReserve sdk.Dec) sdk.Dec {
 	// no need to divide by a, since a = 1.
 	return solveCfmm(baseReserve, quoteReserve, a)
 }
+
+// returns outAmt as a decimal
+func (pa *Pool) calcOutAmtGivenIn(tokenIn sdk.Coin, tokenOutDenom string, swapFee sdk.Dec) (sdk.Dec, error) {
+	reserves, err := pa.getPoolAmts(tokenIn.Denom, tokenOutDenom)
+	if err != nil {
+		return sdk.Dec{}, err
+	}
+	tokenInSupply := reserves[0].ToDec()
+	tokenOutSupply := reserves[1].ToDec()
+	// We are solving for the amount of token out, hence x = tokenOutSupply, y = tokenInSupply
+	outAmt := solveCfmm(tokenOutSupply, tokenInSupply, tokenIn.Amount.ToDec())
+	return outAmt, nil
+}
+
+// returns inAmt as a decimal
+func (pa *Pool) calcInAmtGivenOut(tokenOut sdk.Coin, tokenInDenom string, swapFee sdk.Dec) (sdk.Dec, error) {
+	reserves, err := pa.getPoolAmts(tokenInDenom, tokenOut.Denom)
+	if err != nil {
+		return sdk.Dec{}, err
+	}
+	tokenInSupply := reserves[0].ToDec()
+	tokenOutSupply := reserves[1].ToDec()
+	// 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
+	inAmtRaw := solveCfmm(tokenInSupply, tokenOutSupply, tokenOut.Amount.ToDec().Neg())
+	inAmt := inAmtRaw.NegMut()
+	return inAmt, nil
+}

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

@@ -33,8 +33,6 @@ func TestCFMMInvariant(t *testing.T) {
 			sdk.NewDec(100),
 			sdk.NewDec(100),
 			sdk.NewDec(1000),
-			// returns 87.445364416281417284
-			// should return 99.84973704262359
 		},
 		// {
 		// 	sdk.NewDec(100000),

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

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
 	"github.com/osmosis-labs/osmosis/v7/x/gamm/types"
 )
@@ -68,34 +69,86 @@ func (pa Pool) getPoolAmts(denoms ...string) ([]sdk.Int, error) {
 	return result, nil
 }
 
-// These should all get moved to amm.go
+// updatePoolLiquidityForSwap updates the pool liquidity.
+// It requires caller to validate that tokensIn and tokensOut only consist of
+// denominations in the pool.
+// The function sanity checks this, and panics if not the case.
+func (p *Pool) updatePoolLiquidityForSwap(tokensIn sdk.Coins, tokensOut sdk.Coins) {
+	numTokens := p.PoolLiquidity.Len()
+	// update liquidity
+	p.PoolLiquidity = p.PoolLiquidity.Add(tokensIn...).Sub(tokensOut)
+	// sanity check that no new denoms were added
+	if len(p.PoolLiquidity) != numTokens {
+		panic("updatePoolLiquidityForSwap changed number of tokens in pool")
+	}
+}
+
+// TODO: These should all get moved to amm.go
 func (pa Pool) CalcOutAmtGivenIn(ctx sdk.Context, tokenIn sdk.Coins, tokenOutDenom string, swapFee sdk.Dec) (tokenOut sdk.DecCoin, err error) {
 	if tokenIn.Len() != 1 {
-		return sdk.DecCoin{}, errors.New("asdf")
+		return sdk.DecCoin{}, errors.New("stableswap CalcOutAmtGivenIn: tokenIn is of wrong length")
 	}
-	reserves, err := pa.getPoolAmts(tokenIn[0].Denom, tokenOutDenom)
+	amt, err := pa.calcOutAmtGivenIn(tokenIn[0], tokenOutDenom, swapFee)
 	if err != nil {
 		return sdk.DecCoin{}, err
 	}
-	// document which is x vs y
-	outAmt := solveCfmm(reserves[1].ToDec(), reserves[0].ToDec(), tokenIn[0].Amount.ToDec())
-	return sdk.DecCoin{Denom: tokenOutDenom, Amount: outAmt}, nil
+	return sdk.DecCoin{Denom: tokenOutDenom, Amount: amt}, nil
 }
 
 func (pa *Pool) SwapOutAmtGivenIn(ctx sdk.Context, tokenIn sdk.Coins, tokenOutDenom string, swapFee sdk.Dec) (tokenOut sdk.Coin, err error) {
-	return sdk.Coin{}, types.ErrNotImplemented
+	tokenOutDec, err := pa.CalcOutAmtGivenIn(ctx, tokenIn, tokenOutDenom, swapFee)
+	if err != nil {
+		return sdk.Coin{}, err
+	}
+	// we ignore the decimal component, as token out amount must round down
+	tokenOut, _ = tokenOutDec.TruncateDecimal()
+	if !tokenOut.Amount.IsPositive() {
+		return sdk.Coin{}, sdkerrors.Wrapf(types.ErrInvalidMathApprox, "token amount must be positive")
+	}
+	pa.updatePoolLiquidityForSwap(tokenIn, sdk.NewCoins(tokenOut))
+
+	return tokenOut, nil
 }
 
 func (pa Pool) CalcInAmtGivenOut(ctx sdk.Context, tokenOut sdk.Coins, tokenInDenom string, swapFee sdk.Dec) (tokenIn sdk.DecCoin, err error) {
-	return sdk.DecCoin{}, types.ErrNotImplemented
+	if tokenOut.Len() != 1 {
+		return sdk.DecCoin{}, errors.New("stableswap CalcInAmtGivenOut: tokenOut is of wrong length")
+	}
+	// TODO: Refactor this later to handle scaling factors
+	amt, err := pa.calcInAmtGivenOut(tokenOut[0], tokenInDenom, swapFee)
+	if err != nil {
+		return sdk.DecCoin{}, err
+	}
+	// TODO: Once we make calc in amt given out return a Coin
+	// if tokenInDecimal is non-zero, we add 1 to the tokenInCoin
+	// this is because tokenIn must round up
+	// if tokenInDecimal.Amount.IsPositive() {
+	// 	tokenIn.Amount = tokenIn.Amount.AddRaw(1)
+	// }
+	return sdk.DecCoin{Denom: tokenInDenom, Amount: amt}, nil
 }
 
 func (pa *Pool) SwapInAmtGivenOut(ctx sdk.Context, tokenOut sdk.Coins, tokenInDenom string, swapFee sdk.Dec) (tokenIn sdk.Coin, err error) {
-	return sdk.Coin{}, types.ErrNotImplemented
+	tokenInDec, err := pa.CalcInAmtGivenOut(ctx, tokenOut, tokenInDenom, swapFee)
+	if err != nil {
+		return sdk.Coin{}, err
+	}
+	tokenIn, _ = tokenInDec.TruncateDecimal()
+	if !tokenIn.Amount.IsPositive() {
+		return sdk.Coin{}, sdkerrors.Wrapf(types.ErrInvalidMathApprox, "token amount must be positive")
+	}
+	pa.updatePoolLiquidityForSwap(sdk.NewCoins(tokenIn), tokenOut)
+
+	return tokenIn, nil
 }
 
 func (pa Pool) SpotPrice(ctx sdk.Context, baseAssetDenom string, quoteAssetDenom string) (sdk.Dec, error) {
-	return sdk.Dec{}, types.ErrNotImplemented
+	reserves, err := pa.getPoolAmts(baseAssetDenom, quoteAssetDenom)
+	if err != nil {
+		return sdk.Dec{}, err
+	}
+	// TODO: apply scaling factors here
+	return spotPrice(reserves[0].ToDec(), reserves[1].ToDec()), nil
 }
 
 func (pa Pool) CalcJoinPoolShares(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, newLiquidity sdk.Coins, err error) {