feat(stableswap): implement and test JoinPoolNoSwap and CalcJoinPoolNoSwapShares abstractions

x/gamm/pool-models/internal/cfmm_common/lp.go

@@ -84,8 +84,8 @@ func MaximalExactRatioJoin(p types.PoolI, ctx sdk.Context, tokensIn sdk.Coins) (
 		coinShareRatios[i] = shareRatio
 	}
 
-	if minShareRatio.Equal(sdk.MaxSortableDec) {
-		return numShares, remCoins, errors.New("unexpected error in MaximalExactRatioJoin")
+	if minShareRatio.GTE(sdk.MaxSortableDec) {
+		return numShares, remCoins, types.RatioOfTokensInToExistingLiqExceededError{ActualRatio: minShareRatio}
 	}
 
 	remCoins = sdk.Coins{}

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

@@ -1,8 +1,6 @@
 package stableswap
 
 import (
-	"errors"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
 	"github.com/osmosis-labs/osmosis/v12/osmomath"
@@ -317,24 +315,21 @@ func (p *Pool) calcSingleAssetJoinShares(tokenIn sdk.Coin, swapFee sdk.Dec) (sdk
 	return cfmm_common.BinarySearchSingleAssetJoin(p, sdk.NewCoin(tokenIn.Denom, tokenInAmtAfterFee), poolWithAddedLiquidityAndShares)
 }
 
+// TODO: godoc
 // We can mutate pa here
 // TODO: some day switch this to a COW wrapped pa, for better perf
-func (p *Pool) joinPoolSharesInternal(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, newLiquidity sdk.Coins, err error) {
+func (p *Pool) joinPoolSharesInternal(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (sdk.Int, sdk.Coins, error) {
 	if len(tokensIn) == 1 {
-		numShares, err = p.calcSingleAssetJoinShares(tokensIn[0], swapFee)
-		newLiquidity = tokensIn
-		return numShares, newLiquidity, err
-	} else if len(tokensIn) != p.NumAssets() || !tokensIn.DenomsSubsetOf(p.GetTotalPoolLiquidity(ctx)) {
-		return sdk.ZeroInt(), sdk.NewCoins(), errors.New(
-			"stableswap pool only supports LP'ing with one asset, or all assets in pool")
+		numShares, err := p.calcSingleAssetJoinShares(tokensIn[0], swapFee)
+		p.updatePoolForJoin(tokensIn, numShares)
+		return numShares, tokensIn, err
 	}
-
-	// Add all exact coins we can (no swap). ctx arg doesn't matter for Stableswap
-	numShares, remCoins, err := cfmm_common.MaximalExactRatioJoin(p, sdk.Context{}, tokensIn)
+	numShares, tokensJoined, err := p.joinPoolNoSwapSharesInternal(ctx, tokensIn, swapFee)
 	if err != nil {
-		return sdk.ZeroInt(), sdk.NewCoins(), err
+		return sdk.Int{}, sdk.Coins{}, err
 	}
-	p.updatePoolForJoin(tokensIn.Sub(remCoins), numShares)
+
+	remCoins := tokensIn.Sub(tokensJoined)
 
 	for _, coin := range remCoins {
 		// TODO: Perhaps add a method to skip if this is too small.
@@ -348,3 +343,31 @@ func (p *Pool) joinPoolSharesInternal(ctx sdk.Context, tokensIn sdk.Coins, swapF
 
 	return numShares, tokensIn, nil
 }
+
+// joinPoolNoSwapSharesInternal joins the pool with the given amount of tokens in and swap fee.
+// On success, returns the number of shares and tokens joined.
+// Returns error if:
+// - one token in given
+// - tokens in assets do not match pool assets
+// - equal join fails internally
+func (p *Pool) joinPoolNoSwapSharesInternal(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, tokensJoined sdk.Coins, err error) {
+	poolLiquidity := p.GetTotalPoolLiquidity(ctx)
+	if tokensIn.Len() == 1 {
+		return sdk.Int{}, sdk.Coins{}, types.ErrStableSwapNoSwapJoinNeedsMultiAssetsIn
+	}
+	if !(tokensIn.DenomsSubsetOf(poolLiquidity) && poolLiquidity.DenomsSubsetOf(tokensIn)) {
+		return sdk.Int{}, sdk.Coins{}, types.StableSwapPoolAssetsDoNotEqualTokensInJoinError{PoolAssets: poolLiquidity, TokensIn: tokensIn}
+	}
+
+	// Add all exact coins we can (no swap). ctx arg doesn't matter for Stableswap
+	numShares, remCoins, err := cfmm_common.MaximalExactRatioJoin(p, ctx, tokensIn)
+	if err != nil {
+		return sdk.Int{}, sdk.Coins{}, err
+	}
+
+	tokensJoined = tokensIn.Sub(remCoins)
+
+	p.updatePoolForJoin(tokensJoined, numShares)
+
+	return numShares, tokensJoined, nil
+}

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

@@ -13,6 +13,7 @@ import (
 	"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/pool-models/internal/test_helpers"
+	types "github.com/osmosis-labs/osmosis/v12/x/gamm/types"
 )
 
 // CFMMTestCase defines a testcase for stableswap pools
@@ -24,6 +25,13 @@ type CFMMTestCase struct {
 	expectPanic bool
 }
 
+const (
+	baseAmount = 1000000000000
+	extraDenom = "iamextra"
+	denomA     = "usdc"
+	denomB     = "ist"
+)
+
 var (
 	overflowDec           = osmomath.NewDecFromBigInt(new(big.Int).Sub(new(big.Int).Exp(big.NewInt(2), big.NewInt(1024), nil), big.NewInt(1)))
 	twoAssetCFMMTestCases = map[string]CFMMTestCase{
@@ -371,12 +379,23 @@ var (
 			expectPanic: true,
 		},
 	}
+
+	baseInitialPoolLiquidity = sdk.NewCoins(
+		sdk.NewInt64Coin(denomA, baseAmount),
+		sdk.NewInt64Coin(denomB, baseAmount))
+	tenPercentOfBaseInt  = sdk.NewInt(baseAmount / 10)
+	fivePercentOfBaseInt = sdk.NewInt(baseAmount / 20)
 )
 
 type StableSwapTestSuite struct {
 	test_helpers.CfmmCommonTestSuite
 }
 
+func (suite StableSwapTestSuite) validatePoolLiquidityAndShares(ctx sdk.Context, pool types.PoolI, expectedLiquidty sdk.Coins, expectedShares sdk.Int) {
+	suite.Require().Equal(expectedLiquidty, pool.GetTotalPoolLiquidity(ctx))
+	suite.Require().Equal(expectedShares, pool.GetTotalShares())
+}
+
 func TestStableSwapTestSuite(t *testing.T) {
 	suite.Run(t, new(StableSwapTestSuite))
 }
@@ -570,14 +589,124 @@ 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(), poolLiquidity, swapFeeDec, exitFeeDec)
 				suite.Require().NotNil(pool)
 				test_helpers.TestCalculateAmountOutAndIn_InverseRelationship(suite.T(), ctx, pool, poolLiquidityIn.Denom, poolLiquidityOut.Denom, tc.initialCalcOut, swapFeeDec)
 			})
 		}
 	}
 }
 
+func (suite *StableSwapTestSuite) TestJoinPoolNoSwapSharesInternal() {
+	tests := map[string]struct {
+		initialPoolLiquidity sdk.Coins
+
+		tokensIn sdk.Coins
+		swapFee  sdk.Dec
+
+		expectedNumShares    sdk.Int
+		expectedTokensJoined sdk.Coins
+		expectError          error
+	}{
+		// We consider this test case as base case.
+		// The names of the rest of the test cases only mention changes
+		// relative to this base case.
+		"two-asset; zero fees; equal tokensIn": {
+			initialPoolLiquidity: baseInitialPoolLiquidity,
+
+			// denomA = 10%;. denomB = 10% of initial pool liquidity
+			tokensIn: sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt), sdk.NewCoin(denomB, tenPercentOfBaseInt)),
+			swapFee:  sdk.ZeroDec(),
+
+			expectedNumShares:    types.InitPoolSharesSupply.ToDec().Mul(sdk.NewDecWithPrec(1, 1)).TruncateInt(),
+			expectedTokensJoined: sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt), sdk.NewCoin(denomB, tenPercentOfBaseInt)),
+		},
+		"unequal tokens in, join only equal amounts": {
+			initialPoolLiquidity: baseInitialPoolLiquidity,
+
+			// denomA = 10%;. denomB = 5% of initial pool liquidity
+			tokensIn: sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt), sdk.NewCoin(denomB, fivePercentOfBaseInt)),
+			swapFee:  sdk.ZeroDec(),
+
+			// corresponds to denomB's minimum of tokensIn relative to initial pool liquidity of 5%
+			expectedNumShares:    types.InitPoolSharesSupply.ToDec().Mul(sdk.NewDecWithPrec(5, 2)).TruncateInt(),
+			expectedTokensJoined: sdk.NewCoins(sdk.NewCoin(denomA, fivePercentOfBaseInt), sdk.NewCoin(denomB, fivePercentOfBaseInt)),
+		},
+		"one asset - error": {
+			initialPoolLiquidity: baseInitialPoolLiquidity,
+
+			tokensIn: sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt)),
+
+			expectError: types.ErrStableSwapNoSwapJoinNeedsMultiAssetsIn,
+		},
+		"token in denoms is not subset of pool asset denoms - error": {
+			initialPoolLiquidity: baseInitialPoolLiquidity,
+
+			// proportions are irrelevant here
+			tokensIn: sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt), sdk.NewCoin(extraDenom, tenPercentOfBaseInt)),
+
+			expectError: types.StableSwapPoolAssetsDoNotEqualTokensInJoinError{
+				PoolAssets: baseInitialPoolLiquidity,
+				TokensIn:   sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt), sdk.NewCoin(extraDenom, tenPercentOfBaseInt)),
+			},
+		},
+		"pool assets are not subset of token in denoms - error": {
+			initialPoolLiquidity: baseInitialPoolLiquidity.Add(sdk.NewCoin(extraDenom, sdk.NewInt(baseAmount))),
+
+			// proportions are irrelevant here
+			tokensIn: sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt), sdk.NewCoin(denomB, tenPercentOfBaseInt)),
+
+			expectError: types.StableSwapPoolAssetsDoNotEqualTokensInJoinError{
+				PoolAssets: baseInitialPoolLiquidity.Add(sdk.NewCoin(extraDenom, tenPercentOfBaseInt)),
+				TokensIn:   sdk.NewCoins(sdk.NewCoin(denomA, tenPercentOfBaseInt), sdk.NewCoin(denomB, tenPercentOfBaseInt)),
+			},
+		},
+		"try joinining with amount much larger than existing liquidity": {
+			initialPoolLiquidity: baseInitialPoolLiquidity,
+
+			// We force the amount the amounts of tokens in to be much larger than the supported ratio.
+			// See cfmm_common.MaximalExactRatioJoin(...) for more details.
+			tokensIn: sdk.NewCoins(sdk.NewCoin(denomA, sdk.MaxSortableDec.Add(sdk.OneDec()).MulInt64(baseAmount).TruncateInt()), sdk.NewCoin(denomB, sdk.MaxSortableDec.Add(sdk.OneDec()).MulInt64(baseAmount).TruncateInt())),
+			swapFee:  sdk.ZeroDec(),
+
+			// See cfmm_common.MaximalExactRatioJoin(...) for details about this ratio.
+			expectError: types.RatioOfTokensInToExistingLiqExceededError{ActualRatio: sdk.MaxSortableDec.Add(sdk.OneDec()).MulInt64(baseAmount).TruncateInt().ToDec().QuoInt(sdk.NewInt(baseAmount))},
+		},
+
+		// TODO: multi-asset, non-zero swap fee, non-base amounts.
+	}
+
+	for name, tc := range tests {
+		suite.Run(name, func() {
+			ctx := suite.CreateTestContext()
+
+			poolI := createTestPool(suite.T(), tc.initialPoolLiquidity, tc.swapFee, sdk.ZeroDec())
+
+			pool, ok := (poolI).(*Pool)
+			suite.Require().True(ok)
+
+			numShares, tokensJoined, err := pool.joinPoolNoSwapSharesInternal(ctx, tc.tokensIn, tc.swapFee)
+
+			if tc.expectError != nil {
+				suite.Require().Error(err)
+				suite.Require().Equal(sdk.Int{}, numShares)
+				suite.Require().Equal(sdk.Coins{}, tokensJoined)
+
+				// validate pool is not updated
+				suite.validatePoolLiquidityAndShares(ctx, pool, tc.initialPoolLiquidity, types.InitPoolSharesSupply)
+				return
+			}
+
+			suite.Require().NoError(err)
+			suite.Require().Equal(tc.expectedNumShares, numShares)
+			suite.Require().Equal(tc.expectedTokensJoined, tokensJoined)
+
+			// validate pool is updated
+			suite.validatePoolLiquidityAndShares(ctx, pool, tc.initialPoolLiquidity.Add(tc.expectedTokensJoined...), types.InitPoolSharesSupply.Add(tc.expectedNumShares))
+		})
+	}
+}
+
 func calcUReserve(remReserves []osmomath.BigDec) osmomath.BigDec {
 	uReserve := osmomath.OneDec()
 	for _, assetReserve := range remReserves {

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

@@ -286,19 +286,43 @@ func (p *Pool) CalcJoinPoolShares(ctx sdk.Context, tokensIn sdk.Coins, swapFee s
 	return pCopy.joinPoolSharesInternal(ctx, tokensIn, swapFee)
 }
 
-// TODO: implement this
-func (p *Pool) CalcJoinPoolNoSwapShares(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, newLiquidity sdk.Coins, err error) {
-	return sdk.ZeroInt(), nil, err
+// CalcJoinPoolNoSwapShares creates and returns the number of shares
+// from joining and updating the pool with the given tokensIn and swapFee.
+// Only works more than one token in tokensIn and equal proportions
+// relative to existing liquidity of thereof. Does not mutate pool state.
+// implements Pool interface. See its definition for more details.
+func (p *Pool) CalcJoinPoolNoSwapShares(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, tokensJoined sdk.Coins, err error) {
+	// N.B. we simulate the calculation by attempting to join a copy of the original pool.
+	// The original pool is not modified.
+	pCopy := p.Copy()
+	numShares, tokensJoined, err = pCopy.joinPoolNoSwapSharesInternal(ctx, tokensIn, swapFee)
+	if err != nil {
+		return sdk.Int{}, sdk.Coins{}, err
+	}
+	// Should never happen but we check anyway.
+	if !tokensJoined.IsEqual(tokensIn) {
+		return sdk.Int{}, sdk.Coins{}, types.TokensJoinedDoNotEqualTokensInNoSwapError{NewLiquidity: tokensJoined, TokensIn: tokensIn}
+	}
+	return numShares, tokensJoined, nil
 }
 
-func (p *Pool) JoinPool(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, err error) {
-	numShares, _, err = p.joinPoolSharesInternal(ctx, tokensIn, swapFee)
-	return numShares, err
+func (p *Pool) JoinPool(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numSharesOut sdk.Int, err error) {
+	numSharesOut, _, err = p.joinPoolSharesInternal(ctx, tokensIn, swapFee)
+	return numSharesOut, err
 }
 
-// TODO: implement this
+// JoinPoolNoSwap returns the number of shares that would be created
+// if we were to join the pool with the given tokensIn and swapFee.
+// Only works more than one token in tokensIn and equal proportions
+// relative to existing liquidity of thereof. Mutates pool state.
+// implements Pool interface. See its definition for more details.
 func (p *Pool) JoinPoolNoSwap(ctx sdk.Context, tokensIn sdk.Coins, swapFee sdk.Dec) (numShares sdk.Int, err error) {
-	return sdk.ZeroInt(), err
+	numSharesOut, tokensJoined, err := p.CalcJoinPoolNoSwapShares(ctx, tokensIn, swapFee)
+	if err != nil {
+		return sdk.Int{}, err
+	}
+	p.updatePoolForJoin(tokensJoined, numSharesOut)
+	return numSharesOut, nil
 }
 
 func (p *Pool) ExitPool(ctx sdk.Context, exitingShares sdk.Int, exitFee sdk.Dec) (exitingCoins sdk.Coins, err error) {