From 62273ccfb4421f13517e51adb2e1e5a21ea85e97 Mon Sep 17 00:00:00 2001
From: Minh Nhat Hoang <nhat.hm@outlook.com>
Date: Tue, 17 Dec 2024 11:50:03 +0700
Subject: [PATCH] fix(mx-trading): reduce inventory limit

---
 .../mx-trading/pool_simulator.go              |  6 ++-
 .../mx-trading/pool_simulator_test.go         | 43 +++++++++++++------
 2 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/pkg/liquidity-source/mx-trading/pool_simulator.go b/pkg/liquidity-source/mx-trading/pool_simulator.go
index faa302b8c..7d5cfce24 100644
--- a/pkg/liquidity-source/mx-trading/pool_simulator.go
+++ b/pkg/liquidity-source/mx-trading/pool_simulator.go
@@ -130,7 +130,11 @@ func (p *PoolSimulator) CalculateLimit() map[string]*big.Int {
 	tokens, reserves := p.GetTokens(), p.GetReserves()
 	inventory := make(map[string]*big.Int, len(tokens))
 	for i, token := range tokens {
-		inventory[token] = new(big.Int).Set(reserves[i])
+		var reducedReserve big.Float
+		reducedReserve.SetInt(reserves[i]).Mul(&reducedReserve, big.NewFloat(0.95))
+		// Reduce each token's reserve by 5% to prevent the total quote amount of a token
+		// from potentially exceeding the maker's balance when building the route
+		inventory[token], _ = reducedReserve.Int(nil)
 	}
 
 	return inventory
diff --git a/pkg/liquidity-source/mx-trading/pool_simulator_test.go b/pkg/liquidity-source/mx-trading/pool_simulator_test.go
index 01765b649..33244b955 100644
--- a/pkg/liquidity-source/mx-trading/pool_simulator_test.go
+++ b/pkg/liquidity-source/mx-trading/pool_simulator_test.go
@@ -84,13 +84,19 @@ func TestPoolSimulator_GetAmountOut(t *testing.T) {
 			expectedErr: ErrAmountInIsGreaterThanTotalLevelSize,
 		},
 		{
-			name:              "it should return correct amountOut1 when all levels are filled",
-			amountIn0:         bignumber.NewBig("4929473425227476992"),
+			name:        "it should return correct amountOut1 when all levels are filled",
+			amountIn0:   bignumber.NewBig("4929473425227476992"),
+			expectedErr: ErrAmountOutIsGreaterThanInventory,
+
+			// Expected output before reducing inventory limit
 			expectedAmountOut: bignumber.NewBig("16488225768595991298048"),
 		},
 		{
-			name:              "it should return correct amountOut0 when all levels are filled",
-			amountIn1:         bignumber.NewBig("201363156713348041539584"),
+			name:        "it should return correct amountOut0 when all levels are filled",
+			amountIn1:   bignumber.NewBig("201363156713348041539584"),
+			expectedErr: ErrAmountOutIsGreaterThanInventory,
+
+			// Expected output before reducing inventory limit
 			expectedAmountOut: bignumber.NewBig("59925038314246815744"),
 		},
 		{
@@ -198,8 +204,11 @@ func TestPoolSimulator_UpdateBalance(t *testing.T) {
 			},
 		},
 		{
-			name:                         "fill all levels 0to1",
-			amountIn0:                    bignumber.NewBig("4929473425227476992"),
+			name:        "fill all levels 0to1",
+			amountIn0:   bignumber.NewBig("4929473425227476992"),
+			expectedErr: ErrAmountOutIsGreaterThanInventory,
+
+			// Expected output before reducing inventory limit
 			expectedZeroToOnePriceLevels: nil,
 			expectedOneToZeroPriceLevels: []PriceLevel{
 				{Size: 546.1, Price: 0.00029818453400477844},
@@ -237,8 +246,11 @@ func TestPoolSimulator_UpdateBalance(t *testing.T) {
 			},
 		},
 		{
-			name:      "fill token1",
-			amountIn1: bignumber.NewBig("201362500000000000000000"),
+			name:        "fill token1",
+			amountIn1:   bignumber.NewBig("201362500000000000000000"),
+			expectedErr: ErrAmountOutIsGreaterThanInventory,
+
+			// Expected output before reducing inventory limit
 			expectedZeroToOnePriceLevels: []PriceLevel{
 				{Size: 0.719, Price: 3347.4385889037885},
 				{Size: 0.015, Price: 3347.141106167435},
@@ -284,8 +296,8 @@ func TestPoolSimulator_UpdateBalance(t *testing.T) {
 			}
 			limit := swaplimit.NewInventory("mx-trading", p.CalculateLimit())
 
-			assert.Equal(t, entityPoolData.Reserves[0], limit.GetLimit(p.token0.Address).String())
-			assert.Equal(t, entityPoolData.Reserves[1], limit.GetLimit(p.token1.Address).String())
+			assert.Equal(t, getReducedReserve(p.GetReserves()[0]).String(), limit.GetLimit(p.token0.Address).String())
+			assert.Equal(t, getReducedReserve(p.GetReserves()[1]).String(), limit.GetLimit(p.token1.Address).String())
 
 			calcAmountOutResult, err := p.CalcAmountOut(pool.CalcAmountOutParams{
 				TokenAmountIn: pool.TokenAmount{Token: token, Amount: amountIn},
@@ -309,13 +321,20 @@ func TestPoolSimulator_UpdateBalance(t *testing.T) {
 
 			tokenInIndex := p.GetTokenIndex(token)
 			assert.Equal(t,
-				new(big.Int).Add(p.GetReserves()[tokenInIndex], amountIn).String(),
+				new(big.Int).Add(getReducedReserve(p.GetReserves()[tokenInIndex]), amountIn).String(),
 				limit.GetLimit(token).String(),
 			)
 			assert.Equal(t,
-				new(big.Int).Sub(p.GetReserves()[1-tokenInIndex], calcAmountOutResult.TokenAmountOut.Amount).String(),
+				new(big.Int).Sub(getReducedReserve(p.GetReserves()[1-tokenInIndex]), calcAmountOutResult.TokenAmountOut.Amount).String(),
 				limit.GetLimit(calcAmountOutResult.TokenAmountOut.Token).String(),
 			)
 		})
 	}
 }
+
+func getReducedReserve(reserve *big.Int) *big.Int {
+	reducedReserve := new(big.Float).SetInt(reserve)
+	reducedReserve.Mul(reducedReserve, big.NewFloat(0.95))
+	reserveInt, _ := reducedReserve.Int(nil)
+	return reserveInt
+}