feat: improve math for single asset join

CHANGELOG.md

@@ -41,6 +41,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased 
 
+### Features
+* [#1115](https://github.com/NibiruChain/nibiru/pull/1115) - feat: improve single asset join calculation
+
 ### Bug Fixes
 * [#1113](https://github.com/NibiruChain/nibiru/pull/1113) - fix: fix quick simulation issue
 * [#1114](https://github.com/NibiruChain/nibiru/pull/1114) - fix(dex): fix single asset join

x/common/common.go

@@ -38,7 +38,7 @@ var (
 	// Precision for int representation in sdk.Int objects
 	Precision = int64(1_000_000)
 
-	BigIntPrecision = 18
+	BigIntPrecision = int64(18)
 )
 
 //-----------------------------------------------------------------------------

x/dex/keeper/keeper_test.go

@@ -669,7 +669,6 @@ func TestJoinPoolAllAssets(t *testing.T) {
 		{
 			name: "join with some assets, pool empty in one side, none remaining",
 			joinerInitialFunds: sdk.NewCoins(
-				sdk.NewInt64Coin("bar", 100),
 				sdk.NewInt64Coin("foo", 100),
 			),
 			initialPool: mock.DexPool(
@@ -682,20 +681,20 @@ func TestJoinPoolAllAssets(t *testing.T) {
 			tokensIn: sdk.NewCoins(
 				sdk.NewInt64Coin("foo", 100),
 			),
-			expectedNumSharesOut: sdk.NewInt64Coin(shareDenom, 900),
-			expectedRemCoins:     sdk.NewCoins(sdk.NewInt64Coin("foo", 1)),
+			expectedNumSharesOut: sdk.NewInt64Coin(shareDenom, 904),
+			expectedRemCoins:     sdk.NewCoins(),
 			expectedJoinerFinalFunds: sdk.NewCoins(
-				sdk.NewInt64Coin(shareDenom, 50),
+				sdk.NewInt64Coin(shareDenom, 904),
 				sdk.NewInt64Coin("bar", 0),
 				sdk.NewInt64Coin("foo", 0),
 			),
 			expectedFinalPool: mock.DexPool(
 				/*poolId=*/ 1,
 				/*assets=*/ sdk.NewCoins(
 					sdk.NewInt64Coin("bar", 100),
-					sdk.NewInt64Coin("foo", 100),
+					sdk.NewInt64Coin("foo", 101),
 				),
-				/*shares=*/ 1000),
+				/*shares=*/ 1004),
 		},
 		{
 			name: "join with some assets, but swap done",
@@ -714,20 +713,20 @@ func TestJoinPoolAllAssets(t *testing.T) {
 				sdk.NewInt64Coin("bar", 50),
 				sdk.NewInt64Coin("foo", 75),
 			),
-			expectedNumSharesOut: sdk.NewInt64Coin(shareDenom, 61),
+			expectedNumSharesOut: sdk.NewInt64Coin(shareDenom, 62),
 			expectedRemCoins:     sdk.NewCoins(),
 			expectedJoinerFinalFunds: sdk.NewCoins(
-				sdk.NewInt64Coin(shareDenom, 50),
-				sdk.NewInt64Coin("bar", 35),
-				sdk.NewInt64Coin("foo", 35),
+				sdk.NewInt64Coin(shareDenom, 62),
+				sdk.NewInt64Coin("bar", 50),
+				sdk.NewInt64Coin("foo", 25),
 			),
 			expectedFinalPool: mock.DexPool(
 				/*poolId=*/ 1,
 				/*assets=*/ sdk.NewCoins(
 					sdk.NewInt64Coin("bar", 150),
 					sdk.NewInt64Coin("foo", 175),
 				),
-				/*shares=*/ 161),
+				/*shares=*/ 162),
 		},
 	}
 
@@ -749,6 +748,7 @@ func TestJoinPoolAllAssets(t *testing.T) {
 			require.Equal(t, tc.expectedFinalPool, pool)
 			require.Equal(t, tc.expectedNumSharesOut, numSharesOut)
 			require.Equal(t, tc.expectedRemCoins, remCoins)
+			require.Equal(t, tc.expectedJoinerFinalFunds, app.BankKeeper.GetAllBalances(ctx, joinerAddr))
 		})
 	}
 }

x/dex/types/pool.go

@@ -122,8 +122,7 @@ func (pool *Pool) AddTokensToPool(tokensIn sdk.Coins) (
 
 /*
 Adds tokens to a pool optimizing the amount of shares (swap + join) and updates the pool balances (i.e. liquidity).
-We join with tokens first, and then realize a single asset join by computing the optimal swap amount and then joining
-the pool with the assets.
+We maximally join with both tokens first, and then perform a single asset join with the remaining assets.
 
 This function is only necessary for balancer pool. Stableswap pool already takes all the deposit from the user.
 
@@ -154,54 +153,14 @@ func (pool *Pool) AddAllTokensToPool(tokensIn sdk.Coins) (
 		return
 	}
 
-	swapToken, err := pool.SwapForSwapAndJoin(remCoins[0])
-	if err != nil {
-		return
-	}
-	if swapToken.Amount.LT(sdk.OneInt()) {
-		return pool.AddTokensToPool(tokensIn)
-	}
-
-	index, _, err := pool.getPoolAssetAndIndex(swapToken.Denom)
-
-	if err != nil {
-		return
-	}
-
-	otherDenom := pool.PoolAssets[1-index].Token.Denom
-	tokenOut, err := pool.CalcOutAmtGivenIn(
-		/*tokenIn=*/ swapToken,
-		/*tokenOutDenom=*/ otherDenom,
-		/*noFee=*/ true,
-	)
-
-	if err != nil {
-		return
-	}
-
-	err = pool.ApplySwap(swapToken, tokenOut)
-
+	numShares2nd, _, err := pool.AddTokensToPool(remCoins)
 	if err != nil {
 		return
 	}
 
-	tokensIn = sdk.Coins{
-		{
-			Denom:  swapToken.Denom,
-			Amount: remCoins[0].Amount.Sub(swapToken.Amount),
-		},
-		{
-			Denom:  otherDenom,
-			Amount: tokenOut.Amount,
-		},
-	}.Sort()
-
-	numShares2nd, remCoins2nd, err := pool.AddTokensToPool(tokensIn)
-	if err != nil {
-		return
-	}
-
-	return numShares2nd.Add(numShares), remCoins2nd, err
+	numShares = numShares2nd.Add(numShares)
+	remCoins = sdk.NewCoins()
+	return
 }
 
 /*

x/dex/types/pool_test.go

@@ -366,22 +366,20 @@ func TestJoinPoolAllTokens(t *testing.T) {
 				sdk.NewInt64Coin("aaa", 10),
 				sdk.NewInt64Coin("bbb", 10),
 			),
-			expectedNumShares: sdk.NewInt(6),
-			expectedRemCoins: sdk.NewCoins(
-				sdk.NewInt64Coin("aaa", 1),
-			),
+			expectedNumShares: sdk.NewInt(7),
+			expectedRemCoins:  sdk.NewCoins(),
 			expectedPool: Pool{
 				PoolAssets: []PoolAsset{
 					{
-						Token:  sdk.NewInt64Coin("aaa", 109),
+						Token:  sdk.NewInt64Coin("aaa", 110),
 						Weight: sdk.NewInt(1 << 30),
 					},
 					{
 						Token:  sdk.NewInt64Coin("bbb", 210),
 						Weight: sdk.NewInt(1 << 30),
 					},
 				},
-				TotalShares: sdk.NewInt64Coin("nibiru/pool/1", 106),
+				TotalShares: sdk.NewInt64Coin("nibiru/pool/1", 107),
 				TotalWeight: sdk.NewInt(2 << 30),
 				PoolParams:  PoolParams{PoolType: PoolType_BALANCER, SwapFee: sdk.ZeroDec()},
 			},
@@ -408,13 +406,11 @@ func TestJoinPoolAllTokens(t *testing.T) {
 				sdk.NewInt64Coin("bbb", 1_345), // 0.09580147 % of pool
 			),
 			expectedNumShares: sdk.NewInt(1173),
-			expectedRemCoins: sdk.NewCoins(
-				sdk.NewInt64Coin("aaa", 1),
-			),
+			expectedRemCoins:  sdk.NewCoins(),
 			expectedPool: Pool{
 				PoolAssets: []PoolAsset{
 					{
-						Token:  sdk.NewInt64Coin("aaa", 3_503_437),
+						Token:  sdk.NewInt64Coin("aaa", 3_503_438),
 						Weight: sdk.NewInt(1 << 30),
 					},
 					{
@@ -447,22 +443,20 @@ func TestJoinPoolAllTokens(t *testing.T) {
 			tokensIn: sdk.NewCoins(
 				sdk.NewInt64Coin("aaa", 4_859), // 0.138885 % of pool
 			),
-			expectedNumShares: sdk.NewInt(693),
-			expectedRemCoins: sdk.NewCoins(
-				sdk.NewInt64Coin("aaa", 2),
-			),
+			expectedNumShares: sdk.NewInt(694),
+			expectedRemCoins:  sdk.NewCoins(),
 			expectedPool: Pool{
 				PoolAssets: []PoolAsset{
 					{
-						Token:  sdk.NewInt64Coin("aaa", 3_503_436),
+						Token:  sdk.NewInt64Coin("aaa", 3_503_438),
 						Weight: sdk.NewInt(1 << 30),
 					},
 					{
 						Token:  sdk.NewInt64Coin("bbb", 1_403_945),
 						Weight: sdk.NewInt(1 << 30),
 					},
 				},
-				TotalShares: sdk.NewInt64Coin("nibiru/pool/1", 1_000_693),
+				TotalShares: sdk.NewInt64Coin("nibiru/pool/1", 1_000_694),
 				TotalWeight: sdk.NewInt(2 << 30),
 				PoolParams:  PoolParams{PoolType: PoolType_BALANCER, SwapFee: sdk.ZeroDec()},
 			},

x/dex/types/shares.go

@@ -2,70 +2,10 @@ package types
 
 import (
 	"errors"
-	"math/big"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/holiman/uint256"
-
-	"github.com/NibiruChain/nibiru/x/common"
 )
 
-/*
-For a single asset join, compute the number of token that need to be swapped for an optimal swap and join.
-See https://www.notion.so/nibiru/Single-Asset-Join-Math-2075178cb9684062b9b65ad23ff14417
-
-args:
-  - tokenIn: the token to add to the pool
-
-ret:
-  - out: the tokens to swap before joining the pool
-  - err: error if any
-*/
-func (pool *Pool) SwapForSwapAndJoin(tokenIn sdk.Coin) (
-	out sdk.Coin, err error,
-) {
-	PRECISION := int64(1_000_000)
-
-	mu := sdk.OneDec().Quo(sdk.OneDec().Sub(pool.PoolParams.SwapFee))
-	sigma := (sdk.OneDec().Add(mu))
-
-	mu = mu.MulInt64(PRECISION).MulInt64(PRECISION)
-	sigma = sigma.MulInt64(PRECISION)
-
-	lx := pool.PoolBalances().AmountOfNoDenomValidation(tokenIn.Denom).ToDec()
-
-	lxsigmauint256 := uint256.NewInt()
-	lxsigmauint256.SetFromBig(lx.Mul(sigma).BigInt())
-
-	lxmuuint256 := uint256.NewInt()
-	lxmuuint256.SetFromBig(lx.Mul(mu).BigInt())
-
-	xinuint256 := uint256.NewInt()
-	xinuint256.SetFromBig(tokenIn.Amount.ToDec().BigInt())
-
-	squarable := uint256.NewInt().Add(
-		uint256.NewInt().Mul(
-			lxsigmauint256,
-			lxsigmauint256,
-		),
-		uint256.NewInt().Mul(
-			uint256.NewInt().SetUint64(4),
-			uint256.NewInt().Mul(
-				lxmuuint256,
-				xinuint256,
-			),
-		),
-	)
-
-	BigInt := &big.Int{}
-	sqrt := BigInt.Quo(BigInt.Sqrt(squarable.ToBig()), big.NewInt(PRECISION))
-
-	sqrtFactor := sdk.NewDecFromBigIntWithPrec(sqrt, int64(common.BigIntPrecision))
-
-	amount := (sigma.Mul(lx).MulInt(sdk.NewInt(-1)).QuoInt64(PRECISION).Add(sqrtFactor)).Quo(sdk.MustNewDecFromStr("2"))
-	return sdk.NewCoin(tokenIn.Denom, amount.TruncateInt()), err
-}
-
 /*
 Takes a pool and the amount of tokens desired to add to the pool,
 and calculates the number of pool shares and remaining coins after theoretically
@@ -90,6 +30,24 @@ func (pool Pool) numSharesOutFromTokensIn(tokensIn sdk.Coins) (
 	maxShareRatio := sdk.ZeroDec()
 
 	poolLiquidity := pool.PoolBalances()
+	if len(tokensIn) == 1 {
+		// From balancer whitepaper, for 2 assets with the same weight, the shares issued are:
+		// P_{supply} * (sqrt(1+((1-f/2) * x_{in})/X)-1)
+
+		one := sdk.OneDec()
+
+		joinShare := tokensIn[0].Amount.ToDec().Mul(one.Sub(pool.PoolParams.SwapFee.Quo(sdk.NewDec(2)))).QuoInt(
+			poolLiquidity.AmountOfNoDenomValidation(tokensIn[0].Denom),
+		).Add(one)
+
+		joinShare, err = joinShare.ApproxSqrt()
+		if err != nil {
+			return
+		}
+
+		numShares = joinShare.Sub(one).MulInt(pool.TotalShares.Amount).TruncateInt()
+		return
+	}
 
 	for i, coin := range tokensIn {
 		shareRatio := coin.Amount.ToDec().QuoInt(