[x/gamm][stableswap]: Add inverse join/exit tests, fix single asset join bug, and remove uneven ratio joins

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

@@ -382,7 +382,18 @@ func (p *Pool) joinPoolSharesInternal(ctx sdk.Context, tokensIn sdk.Coins, swapF
 	}
 	if len(tokensIn) == 1 && tokensIn[0].Amount.GT(sdk.OneInt()) {
 		numShares, err = p.calcSingleAssetJoinShares(tokensIn[0], swapFee)
+		if err != nil {
+			return sdk.ZeroInt(), sdk.NewCoins(), err
+		}
+
 		newLiquidity = tokensIn
+
+		p.updatePoolForJoin(newLiquidity, numShares)
+
+		if err = validatePoolAssets(p.PoolLiquidity, p.ScalingFactor); err != nil {
+			return sdk.ZeroInt(), sdk.NewCoins(), err
+		}
+
 		return numShares, newLiquidity, err
 	} else if len(tokensIn) != p.NumAssets() {
 		return sdk.ZeroInt(), sdk.NewCoins(), errors.New(
@@ -396,20 +407,7 @@ func (p *Pool) joinPoolSharesInternal(ctx sdk.Context, tokensIn sdk.Coins, swapF
 	}
 	p.updatePoolForJoin(tokensIn.Sub(remCoins), numShares)
 
-	tokensJoined := tokensIn
-	for _, coin := range remCoins {
-		// TODO: Perhaps add a method to skip if this is too small.
-		if coin.Amount.GT(sdk.OneInt()) {
-			newShare, err := p.calcSingleAssetJoinShares(coin, swapFee)
-			if err != nil {
-				return sdk.ZeroInt(), sdk.NewCoins(), err
-			}
-			p.updatePoolForJoin(sdk.NewCoins(coin), newShare)
-			numShares = numShares.Add(newShare)
-		} else {
-			tokensJoined = tokensJoined.Sub(sdk.NewCoins(coin))
-		}
-	}
+	tokensJoined := tokensIn.Sub(remCoins)
 
 	if err = validatePoolAssets(p.PoolLiquidity, p.ScalingFactor); err != nil {
 		return sdk.ZeroInt(), sdk.NewCoins(), err

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

@@ -866,9 +866,9 @@ func TestJoinPoolSharesInternal(t *testing.T) {
 			poolAssets:      twoEvenStablePoolAssets,
 			scalingFactors:  defaultTwoAssetScalingFactors,
 			swapFee:         sdk.ZeroDec(),
-			expNumShare:     sdk.NewIntFromUint64(10000000500000000000),
-			expTokensJoined: sdk.NewCoins(sdk.NewCoin("foo", sdk.NewInt(tenPercentOfTwoPoolRaw)), sdk.NewCoin("bar", sdk.NewInt(10+tenPercentOfTwoPoolRaw))),
-			expPoolAssets:   twoAssetPlusTenPercent.Add(sdk.NewCoin("bar", sdk.NewInt(10))),
+			expNumShare:     sdk.NewIntFromUint64(10000000000000000000),
+			expTokensJoined: sdk.NewCoins(sdk.NewCoin("foo", sdk.NewInt(tenPercentOfTwoPoolRaw)), sdk.NewCoin("bar", sdk.NewInt(tenPercentOfTwoPoolRaw))),
+			expPoolAssets:   twoAssetPlusTenPercent,
 			expectPass:      true,
 		},
 		"all-asset pool join attempt exceeds max scaled asset amount": {

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

@@ -10,6 +10,7 @@ import (
 
 	"github.com/osmosis-labs/osmosis/v12/app/apptesting/osmoassert"
 	"github.com/osmosis-labs/osmosis/v12/osmomath"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/internal/cfmm_common"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/types"
 )
 
@@ -639,3 +640,155 @@ func TestSwapInAmtGivenOut(t *testing.T) {
 		})
 	}
 }
+
+func TestInverseJoinPoolExitPool(t *testing.T) {
+	hundredFoo := sdk.NewCoins(sdk.NewCoin("foo", sdk.NewInt(100)))
+	thousandAssetA := sdk.NewCoins(sdk.NewCoin("asset/a", sdk.NewInt(1000)))
+	tenPercentOfTwoPoolRaw := int64(1000000000 / 10)
+	tenPercentOfTwoPoolCoins := sdk.NewCoins(sdk.NewCoin("foo", sdk.NewInt(int64(1000000000/10))), sdk.NewCoin("bar", sdk.NewInt(int64(1000000000/10))))
+	// twoAssetPlusTenPercent := twoEvenStablePoolAssets.Add(tenPercentOfTwoPoolCoins...)
+	type testcase struct {
+		tokensIn           sdk.Coins
+		poolAssets         sdk.Coins
+		unevenJoinedTokens sdk.Coins
+		scalingFactors     []uint64
+		swapFee            sdk.Dec
+		expectPass         bool
+	}
+
+	tests := map[string]testcase{
+		"[single asset join] even two asset pool, no swap fee": {
+			tokensIn:       hundredFoo,
+			poolAssets:     twoEvenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        sdk.ZeroDec(),
+			expectPass:     true,
+		},
+		"[single asset join] uneven two asset pool, no swap fee": {
+			tokensIn:       hundredFoo,
+			poolAssets:     twoUnevenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        sdk.ZeroDec(),
+			expectPass:     true,
+		},
+		"[single asset join] even 3-asset pool, no swap fee": {
+			tokensIn:       thousandAssetA,
+			poolAssets:     threeEvenStablePoolAssets,
+			scalingFactors: defaultThreeAssetScalingFactors,
+			swapFee:        sdk.ZeroDec(),
+			expectPass:     true,
+		},
+		"[single asset join] uneven 3-asset pool, no swap fee": {
+			tokensIn:       thousandAssetA,
+			poolAssets:     threeUnevenStablePoolAssets,
+			scalingFactors: defaultThreeAssetScalingFactors,
+			swapFee:        sdk.ZeroDec(),
+			expectPass:     true,
+		},
+		"[single asset join] even two asset pool, default swap fee": {
+			tokensIn:       hundredFoo,
+			poolAssets:     twoEvenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        defaultSwapFee,
+			expectPass:     true,
+		},
+		"[single asset join] uneven two asset pool, default swap fee": {
+			tokensIn:       hundredFoo,
+			poolAssets:     twoUnevenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        defaultSwapFee,
+			expectPass:     true,
+		},
+		"[single asset join] even 3-asset pool, default swap fee": {
+			tokensIn:       thousandAssetA,
+			poolAssets:     threeEvenStablePoolAssets,
+			scalingFactors: defaultThreeAssetScalingFactors,
+			swapFee:        defaultSwapFee,
+			expectPass:     true,
+		},
+		"[single asset join] uneven 3-asset pool, default swap fee": {
+			tokensIn:       thousandAssetA,
+			poolAssets:     threeUnevenStablePoolAssets,
+			scalingFactors: defaultThreeAssetScalingFactors,
+			swapFee:        defaultSwapFee,
+			expectPass:     true,
+		},
+		"[single asset join] even 3-asset pool, 0.03 swap fee": {
+			tokensIn:       thousandAssetA,
+			poolAssets:     threeEvenStablePoolAssets,
+			scalingFactors: defaultThreeAssetScalingFactors,
+			swapFee:        sdk.MustNewDecFromStr("0.03"),
+			expectPass:     true,
+		},
+		"[single asset join] uneven 3-asset pool, 0.03 swap fee": {
+			tokensIn:       thousandAssetA,
+			poolAssets:     threeUnevenStablePoolAssets,
+			scalingFactors: defaultThreeAssetScalingFactors,
+			swapFee:        sdk.MustNewDecFromStr("0.03"),
+			expectPass:     true,
+		},
+
+		"[all asset join] even two asset pool, same tokenIn ratio": {
+			tokensIn:       tenPercentOfTwoPoolCoins,
+			poolAssets:     twoEvenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        sdk.ZeroDec(),
+			expectPass:     true,
+		},
+		"[all asset join] even two asset pool, different tokenIn ratio with pool": {
+			tokensIn:       sdk.NewCoins(sdk.NewCoin("foo", sdk.NewInt(tenPercentOfTwoPoolRaw)), sdk.NewCoin("bar", sdk.NewInt(10+tenPercentOfTwoPoolRaw))),
+			poolAssets:     twoEvenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        sdk.ZeroDec(),
+			expectPass:     true,
+		},
+		"[all asset join] even two asset pool, different tokenIn ratio with pool, nonzero swap fee": {
+			tokensIn:       sdk.NewCoins(sdk.NewCoin("foo", sdk.NewInt(tenPercentOfTwoPoolRaw)), sdk.NewCoin("bar", sdk.NewInt(10+tenPercentOfTwoPoolRaw))),
+			poolAssets:     twoEvenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        defaultSwapFee,
+			expectPass:     true,
+		},
+		"[all asset join] even two asset pool, no tokens in": {
+			tokensIn:       sdk.NewCoins(),
+			poolAssets:     twoEvenStablePoolAssets,
+			scalingFactors: defaultTwoAssetScalingFactors,
+			swapFee:        sdk.ZeroDec(),
+			expectPass:     true,
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			ctx := sdk.Context{}
+			p := poolStructFromAssets(tc.poolAssets, tc.scalingFactors)
+
+			// we join then exit the pool
+			shares, err := p.JoinPool(ctx, tc.tokensIn, tc.swapFee)
+			tokenOut, err := p.ExitPool(ctx, shares, defaultExitFee)
+
+			// if single asset join, we swap output tokens to input denom to test the full inverse relationship
+			if len(tc.tokensIn) == 1 {
+				tokenOutAmt, err := cfmm_common.SwapAllCoinsToSingleAsset(&p, ctx, tokenOut, tc.tokensIn[0].Denom)
+				require.NoError(t, err)
+				tokenOut = sdk.NewCoins(sdk.NewCoin(tc.tokensIn[0].Denom, tokenOutAmt))
+			}
+
+			// if single asset join, we expect output token swapped into the input denom to be input minus swap fee
+			var expectedTokenOut sdk.Coins
+			if len(tc.tokensIn) == 1 {
+				expectedAmt := (tc.tokensIn[0].Amount.ToDec().Mul(sdk.OneDec().Sub(tc.swapFee))).TruncateInt()
+				expectedTokenOut = sdk.NewCoins(sdk.NewCoin(tc.tokensIn[0].Denom, expectedAmt))
+			} else {
+				expectedTokenOut = tc.tokensIn
+			}
+
+			if tc.expectPass {
+				finalPoolLiquidity := p.GetTotalPoolLiquidity(ctx)
+				require.True(t, tokenOut.IsAllLTE(expectedTokenOut))
+				require.True(t, finalPoolLiquidity.IsAllGTE(tc.poolAssets))
+			}
+			osmoassert.ConditionalError(t, !tc.expectPass, err)
+		})
+	}
+}