Add mutative, non mutative method for PoolI (#1147)

* Add mutative, non mutative method for PoolI

* Remove ApplySwap and its usage

* Add internal steps for liquidity update before swap calc

x/gamm/keeper/swap.go

@@ -47,15 +47,13 @@ func (k Keeper) swapExactAmountIn(
 	if tokenIn.Denom == tokenOutDenom {
 		return sdk.Int{}, errors.New("cannot trade same denomination in and out")
 	}
-
 	tokensIn := sdk.Coins{tokenIn}
 
-	tokenOutDecCoin, err := pool.CalcOutAmtGivenIn(ctx, tokensIn, tokenOutDenom, swapFee)
+	tokenOutCoin, err := pool.SwapOutAmtGivenIn(ctx, tokensIn, tokenOutDenom, swapFee)
 	if err != nil {
 		return sdk.Int{}, err
 	}
 
-	tokenOutCoin, _ := tokenOutDecCoin.TruncateDecimal()
 	tokenOutAmount = tokenOutCoin.Amount
 
 	if !tokenOutAmount.IsPositive() {
@@ -111,12 +109,10 @@ func (k Keeper) swapExactAmountOut(
 		return sdk.Int{}, sdkerrors.Wrapf(types.ErrTooManyTokensOut,
 			"can't get more tokens out than there are tokens in the pool")
 	}
-
-	tokenInDecCoin, err := pool.CalcInAmtGivenOut(ctx, sdk.Coins{tokenOut}, tokenInDenom, swapFee)
+	tokenInCoin, err := pool.SwapInAmtGivenOut(ctx, sdk.Coins{tokenOut}, tokenInDenom, swapFee)
 	if err != nil {
 		return sdk.Int{}, err
 	}
-	tokenInCoin, _ := tokenInDecCoin.TruncateDecimal()
 	tokenInAmount = tokenInCoin.Amount
 
 	if tokenInAmount.LTE(sdk.ZeroInt()) {
@@ -149,11 +145,7 @@ func (k Keeper) updatePoolForSwap(
 	tokensIn := sdk.Coins{tokenIn}
 	tokensOut := sdk.Coins{tokenOut}
 
-	err := pool.ApplySwap(ctx, tokensIn, tokensOut)
-	if err != nil {
-		return err
-	}
-	err = k.SetPool(ctx, pool)
+	err := k.SetPool(ctx, pool)
 	if err != nil {
 		return err
 	}

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

@@ -4,6 +4,7 @@ import (
 	"errors"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 
 	"github.com/osmosis-labs/osmosis/v7/osmomath"
 	"github.com/osmosis-labs/osmosis/v7/x/gamm/types"
@@ -69,7 +70,32 @@ func (p Pool) CalcOutAmtGivenIn(
 	return sdk.NewDecCoinFromDec(tokenOutDenom, tokenAmountOut), nil
 }
 
-// calcInAmtGivenOut calculates token to be provided, fee added,
+// SwapOutAmtGivenIn is a mutative method for CalcOutAmtGivenIn, which includes the actual swap.
+func (p *Pool) SwapOutAmtGivenIn(
+	ctx sdk.Context,
+	tokensIn sdk.Coins,
+	tokenOutDenom string,
+	swapFee sdk.Dec,
+) (
+	tokenOut sdk.Coin, err error,
+) {
+	tokenOutDecCoin, err := p.CalcOutAmtGivenIn(ctx, tokensIn, tokenOutDenom, swapFee)
+	if err != nil {
+		return sdk.Coin{}, err
+	}
+	tokenOutCoin, _ := tokenOutDecCoin.TruncateDecimal()
+	if !tokenOutCoin.Amount.IsPositive() {
+		return sdk.Coin{}, sdkerrors.Wrapf(types.ErrInvalidMathApprox, "token amount must be positive")
+	}
+
+	err = p.applySwap(ctx, tokensIn, sdk.Coins{tokenOutCoin})
+	if err != nil {
+		return sdk.Coin{}, err
+	}
+	return tokenOutCoin, nil
+}
+
+// CalcInAmtGivenOut calculates token to be provided, fee added,
 // given the swapped out amount, using solveConstantFunctionInvariant.
 func (p Pool) CalcInAmtGivenOut(
 	ctx sdk.Context, tokensOut sdk.Coins, tokenInDenom string, swapFee sdk.Dec) (
@@ -96,8 +122,29 @@ func (p Pool) CalcInAmtGivenOut(
 	return sdk.NewDecCoinFromDec(tokenInDenom, tokenAmountInBeforeFee), nil
 }
 
+// SwapInAmtGivenOut is a mutative method for CalcOutAmtGivenIn, which includes the actual swap.
+func (p *Pool) SwapInAmtGivenOut(
+	ctx sdk.Context, tokensOut sdk.Coins, tokenInDenom string, swapFee sdk.Dec) (
+	tokenIn sdk.Coin, err error,
+) {
+	tokenInDecCoin, err := p.CalcInAmtGivenOut(ctx, tokensOut, tokenInDenom, swapFee)
+	if err != nil {
+		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")
+	}
+
+	err = p.applySwap(ctx, sdk.Coins{tokenInCoin}, tokensOut)
+	if err != nil {
+		return sdk.Coin{}, err
+	}
+	return tokenInCoin, nil
+}
+
 // ApplySwap.
-func (p *Pool) ApplySwap(ctx sdk.Context, tokensIn sdk.Coins, tokensOut sdk.Coins) error {
+func (p *Pool) applySwap(ctx sdk.Context, tokensIn sdk.Coins, tokensOut sdk.Coins) error {
 	// Also ensures that len(tokensIn) = 1 = len(tokensOut)
 	inPoolAsset, outPoolAsset, err := p.parsePoolAssetsCoins(tokensIn, tokensOut)
 	if err != nil {
@@ -173,11 +220,7 @@ func calcPoolOutGivenSingleIn(
 }
 
 // calcPoolOutGivenSingleIn - balance pAo.
-func (p *Pool) singleAssetJoin(tokenIn sdk.Coin, swapFee sdk.Dec) (numShares sdk.Int, err error) {
-	tokenInPoolAsset, err := p.GetPoolAsset(tokenIn.Denom)
-	if err != nil {
-		return sdk.ZeroInt(), err
-	}
+func (p *Pool) calcSingleAssetJoin(tokenIn sdk.Coin, swapFee sdk.Dec, tokenInPoolAsset PoolAsset, totalShares sdk.Int) (numShares sdk.Int, err error) {
 	totalWeight := p.GetTotalWeight()
 	if totalWeight.IsZero() {
 		return sdk.ZeroInt(), errors.New("pool misconfigured, total weight = 0")
@@ -186,7 +229,7 @@ func (p *Pool) singleAssetJoin(tokenIn sdk.Coin, swapFee sdk.Dec) (numShares sdk
 	return calcPoolOutGivenSingleIn(
 		tokenInPoolAsset.Token.Amount.ToDec(),
 		normalizedWeight,
-		p.GetTotalShares().ToDec(),
+		totalShares.ToDec(),
 		tokenIn.Amount.ToDec(),
 		swapFee,
 	).TruncateInt(), nil
@@ -235,35 +278,77 @@ func (p *Pool) maximalExactRatioJoin(tokensIn sdk.Coins) (numShares sdk.Int, rem
 }
 
 func (p *Pool) JoinPool(_ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, err error) {
+	numShares, newLiquidity, err := p.CalcJoinPoolShares(_ctx, tokensIn, swapFee)
+	if err != nil {
+		return sdk.Int{}, err
+	}
+	p.updateLiquidity(numShares, newLiquidity)
+	return numShares, nil
+}
+
+func (p *Pool) CalcJoinPoolShares(_ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, newLiquidity sdk.Coins, err error) {
+	poolAssets := p.GetAllPoolAssets()
+	poolAssetsByDenom := make(map[string]PoolAsset)
+	for _, poolAsset := range poolAssets {
+		poolAssetsByDenom[poolAsset.Token.Denom] = poolAsset
+	}
+	totalShares := p.GetTotalShares()
+
 	if tokensIn.Len() == 1 {
-		numShares, err = p.singleAssetJoin(tokensIn[0], swapFee)
-		p.updateLiquidity(numShares, tokensIn)
-		return numShares, err
+		numShares, err = p.calcSingleAssetJoin(tokensIn[0], swapFee, poolAssetsByDenom[tokensIn[0].Denom], totalShares)
+		newLiquidity = tokensIn
+		return numShares, newLiquidity, err
 	} else if tokensIn.Len() != p.NumAssets() {
-		return sdk.ZeroInt(), errors.New(
+		return sdk.ZeroInt(), sdk.NewCoins(), errors.New(
 			"balancer pool only supports LP'ing with one asset, or all assets in pool")
 	}
 	// Add all exact coins we can (no swap)
 	numShares, remCoins, err := p.maximalExactRatioJoin(tokensIn)
 	if err != nil {
-		return sdk.ZeroInt(), err
+		return sdk.ZeroInt(), sdk.NewCoins(), err
 	}
-	p.updateLiquidity(numShares, tokensIn.Sub(remCoins))
+	// update liquidity for accurate calcSingleAssetJoin calculation
+	newLiquidity = tokensIn.Sub(remCoins)
+	for _, coin := range newLiquidity {
+		poolAsset := poolAssetsByDenom[coin.Denom]
+		poolAsset.Token.Amount = poolAssetsByDenom[coin.Denom].Token.Amount.Add(coin.Amount)
+		poolAssetsByDenom[coin.Denom] = poolAsset
+	}
+	totalShares = totalShares.Add(numShares)
+
 	// if there are coins that couldn't be perfectly joined, do single asset joins for each of them.
 	if !remCoins.Empty() {
 		for _, coin := range remCoins {
-			newShares, err := p.singleAssetJoin(coin, swapFee)
+			newShares, err := p.calcSingleAssetJoin(coin, swapFee, poolAssetsByDenom[coin.Denom], totalShares)
 			if err != nil {
-				return sdk.ZeroInt(), err
+				return sdk.ZeroInt(), sdk.NewCoins(), err
 			}
-			p.updateLiquidity(newShares, sdk.Coins{coin})
+			newLiquidity = newLiquidity.Add(coin)
 			numShares = numShares.Add(newShares)
 		}
 	}
-	return numShares, nil
+	return numShares, newLiquidity, nil
 }
 
 func (p *Pool) ExitPool(ctx sdk.Context, exitingShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error) {
+	exitedCoins, err = p.CalcExitPoolShares(ctx, exitingShares, exitFee)
+	if err != nil {
+		return sdk.Coins{}, err
+	}
+
+	balances := p.GetTotalPoolLiquidity(ctx).Sub(exitedCoins)
+	err = p.UpdatePoolAssetBalances(balances)
+	if err != nil {
+		return sdk.Coins{}, err
+	}
+
+	totalShares := p.GetTotalShares()
+	p.TotalShares = sdk.NewCoin(p.TotalShares.Denom, totalShares.Sub(exitingShares))
+
+	return exitedCoins, nil
+}
+
+func (p *Pool) CalcExitPoolShares(ctx sdk.Context, exitingShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error) {
 	totalShares := p.GetTotalShares()
 	if exitingShares.GTE(totalShares) {
 		return sdk.Coins{}, errors.New(("too many shares out"))
@@ -287,14 +372,6 @@ func (p *Pool) ExitPool(ctx sdk.Context, exitingShares sdk.Int, exitFee sdk.Dec)
 			continue
 		}
 		exitedCoins = exitedCoins.Add(sdk.NewCoin(asset.Denom, exitAmt))
-		// update pool assets for this exit amount.
-		newAmt := asset.Amount.Sub(exitAmt)
-		err = p.UpdatePoolAssetBalance(sdk.NewCoin(asset.Denom, newAmt))
-		if err != nil {
-			return sdk.Coins{}, err
-		}
 	}
-
-	p.TotalShares = sdk.NewCoin(p.TotalShares.Denom, totalShares.Sub(exitingShares))
 	return exitedCoins, nil
 }

x/gamm/types/pool.go

@@ -28,21 +28,22 @@ type PoolI interface {
 	GetTotalShares() sdk.Int
 
 	CalcOutAmtGivenIn(ctx sdk.Context, tokenIn sdk.Coins, tokenOutDenom string, swapFee sdk.Dec) (tokenOut sdk.DecCoin, err error)
-	CalcInAmtGivenOut(ctx sdk.Context, tokenOut sdk.Coins, tokenInDenom string, swapFee sdk.Dec) (tokenIn sdk.DecCoin, err error)
+	SwapOutAmtGivenIn(ctx sdk.Context, tokenIn sdk.Coins, tokenOutDenom string, swapFee sdk.Dec) (tokenOut sdk.Coin, err error)
 
-	// TODO: Ensure this can only be called via gamm
-	// TODO: Think through the API guarantees this is providing in conjunction with the caller being
-	// expected to Set the pool into state as well.
-	ApplySwap(ctx sdk.Context, tokenIn sdk.Coins, tokenOut sdk.Coins) error
+	CalcInAmtGivenOut(ctx sdk.Context, tokenOut sdk.Coins, tokenInDenom string, swapFee sdk.Dec) (tokenIn sdk.DecCoin, err error)
+	SwapInAmtGivenOut(ctx sdk.Context, tokenOut sdk.Coins, tokenInDenom string, swapFee sdk.Dec) (tokenIn sdk.Coin, err error)
 
 	SpotPrice(ctx sdk.Context, baseAssetDenom string, quoteAssetDenom string) (sdk.Dec, error)
 
 	// JoinPool joins the pool, and uses all of the tokensIn provided.
 	// The AMM swaps to whatever the ratio should be and returns the number of shares created.
 	// Internally the pool updates its count for the number of shares in this function.
 	// If the function errors, or should not be mutative, then state must be reverted after this call.
+	CalcJoinPoolShares(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, newLiquidity sdk.Coins, err error)
 	JoinPool(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, err error)
+
 	ExitPool(ctx sdk.Context, numShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error)
+	CalcExitPoolShares(ctx sdk.Context, numShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error)
 }
 
 func NewPoolAddress(poolId uint64) sdk.AccAddress {