refactor/optimization (#655)

* refactor: optimize generic-simple-rate * ft: implement limitorder CloneState * ft: implement hashflow-v3 CloneState * ft: implement saddle CloneState
KyberNetwork · Dec 17, 2024 · 4ab435e · 4ab435e
1 parent 910aed6
commit 4ab435e
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 23 deletions.
diff --git a/pkg/liquidity-source/generic-simple-rate/pool_simulator.go b/pkg/liquidity-source/generic-simple-rate/pool_simulator.go
@@ -26,6 +26,7 @@ type PoolSimulator struct {
 
 var (
 	ErrPoolPaused = errors.New("pool is paused")
+	ErrOverflow   = errors.New("overflow")
 )
 
 func NewPoolSimulator(entityPool entity.Pool) (*PoolSimulator, error) {
@@ -83,8 +84,10 @@ func (p *PoolSimulator) CalcAmountOut(param pool.CalcAmountOutParams) (*pool.Cal
 		return &pool.CalcAmountOutResult{}, fmt.Errorf("tokenInIndex: %v or tokenOutIndex: %v is not correct", tokenInIndex, tokenOutIndex)
 	}
 
-	amountOut := p.calcAmountOut(tokenInIndex, tokenAmountIn.Amount)
-	totalGas := p.gas
+	amountOut, err := p.calcAmountOut(tokenInIndex, tokenAmountIn.Amount)
+	if err != nil {
+		return nil, err
+	}
 
 	return &pool.CalcAmountOutResult{
 		TokenAmountOut: &pool.TokenAmount{
@@ -95,7 +98,7 @@ func (p *PoolSimulator) CalcAmountOut(param pool.CalcAmountOutParams) (*pool.Cal
 			Token:  tokenAmountIn.Token,
 			Amount: bignumber.ZeroBI,
 		},
-		Gas: totalGas,
+		Gas: p.gas,
 	}, nil
 }
 
@@ -126,12 +129,15 @@ func (p *PoolSimulator) GetMetaInfo(_ string, _ string) interface{} {
 	return nil
 }
 
-func (p *PoolSimulator) calcAmountOut(tokenInIndex int, amountIn *big.Int) *uint256.Int {
-	amountOut := new(uint256.Int).Set(uint256.MustFromBig(amountIn))
-	if (p.isRateInversed && tokenInIndex == 0) || (!p.isRateInversed && tokenInIndex == 1) {
+func (p *PoolSimulator) calcAmountOut(tokenInIndex int, amountIn *big.Int) (*uint256.Int, error) {
+	amountOut := new(uint256.Int)
+	if amountOut.SetFromBig(amountIn) {
+		return nil, ErrOverflow
+	}
+	if p.isRateInversed == (tokenInIndex == 0) {
 		amountOut = amountOut.Mul(amountOut, p.rateUnit).Div(amountOut, p.rate)
 	} else {
 		amountOut = amountOut.Div(amountOut, p.rateUnit).Mul(amountOut, p.rate)
 	}
-	return amountOut
+	return amountOut, nil
 }
diff --git a/pkg/liquidity-source/hashflow-v3/pool_simulator.go b/pkg/liquidity-source/hashflow-v3/pool_simulator.go
@@ -159,6 +159,19 @@ func (p *PoolSimulator) CalcAmountIn(params pool.CalcAmountInParams) (*pool.Calc
 	}
 }
 
+func (p *PoolSimulator) CloneState() pool.IPoolSimulator {
+	cloned := *p
+	cloned.ZeroToOnePriceLevels = lo.Map(p.ZeroToOnePriceLevels, func(v PriceLevel, i int) PriceLevel {
+		v.Quote = new(big.Float).Set(v.Quote)
+		return v
+	})
+	cloned.OneToZeroPriceLevels = lo.Map(p.OneToZeroPriceLevels, func(v PriceLevel, i int) PriceLevel {
+		v.Quote = new(big.Float).Set(v.Quote)
+		return v
+	})
+	return &cloned
+}
+
 func (p *PoolSimulator) UpdateBalance(params pool.UpdateBalanceParams) {
 	var amountInAfterDecimals, decimalsPow, amountInBF big.Float
 	amountInBF.SetInt(params.TokenAmountIn.Amount)

diff --git a/pkg/liquidity-source/hashflow-v3/pool_simulator_test.go b/pkg/liquidity-source/hashflow-v3/pool_simulator_test.go
@@ -5,10 +5,11 @@ import (
 	"math/big"
 	"testing"
 
-	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/entity"
-	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/source/pool"
 	"github.com/shopspring/decimal"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/entity"
+	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/source/pool"
 )
 
 var (
@@ -83,11 +84,24 @@ func TestPoolSimulator_GetAmountOut(t *testing.T) {
 				TokenOut: tokenOMG.Address,
 			}
 
-			result, err := poolSimulator.CalcAmountOut(params)
+			sim := poolSimulator.CloneState()
+			cloned := sim.CloneState()
+			result, err := sim.CalcAmountOut(params)
 			assert.Equal(t, tc.expectedErr, err)
-			if tc.expectedErr == nil {
-				assert.Equal(t, 0, result.TokenAmountOut.Amount.Cmp(tc.expectedAmountOut))
+			if tc.expectedErr != nil {
+				return
 			}
+			assert.Equal(t, 0, result.TokenAmountOut.Amount.Cmp(tc.expectedAmountOut))
+
+			sim.UpdateBalance(pool.UpdateBalanceParams{
+				TokenAmountIn:  params.TokenAmountIn,
+				TokenAmountOut: *result.TokenAmountOut,
+				SwapInfo:       result.SwapInfo,
+			})
+
+			clonedRes, err := cloned.CalcAmountOut(params)
+			assert.Equal(t, tc.expectedErr, err)
+			assert.Equal(t, result, clonedRes)
 		})
 	}
 }

diff --git a/pkg/source/limitorder/pool_simulator.go b/pkg/source/limitorder/pool_simulator.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/KyberNetwork/logger"
 	"github.com/goccy/go-json"
+	"github.com/samber/lo"
 
 	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/entity"
 	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/source/pool"
@@ -107,6 +108,20 @@ func (p *PoolSimulator) CalcAmountOut(param pool.CalcAmountOutParams) (*pool.Cal
 	return p.calcAmountOut(param.TokenAmountIn, param.TokenOut, param.Limit)
 }
 
+func (p *PoolSimulator) CloneState() pool.IPoolSimulator {
+	cloned := *p
+	cloned.ordersMapping = lo.MapEntries(p.ordersMapping, func(k int64, v *order) (int64, *order) {
+		c := *v
+		c.FilledTakingAmount = new(big.Int).Set(v.FilledTakingAmount)
+		c.FilledMakingAmount = new(big.Int).Set(v.FilledMakingAmount)
+		if c.AvailableMakingAmount != nil {
+			c.AvailableMakingAmount = new(big.Int).Set(v.AvailableMakingAmount)
+		}
+		return k, &c
+	})
+	return &cloned
+}
+
 func (p *PoolSimulator) UpdateBalance(params pool.UpdateBalanceParams) {
 	swapInfo, ok := params.SwapInfo.(SwapInfo)
 	if !ok {

diff --git a/pkg/source/limitorder/pool_simulator_test.go b/pkg/source/limitorder/pool_simulator_test.go
@@ -1060,14 +1060,16 @@ func TestPool_UpdateBalance(t *testing.T) {
 		limit := swaplimit.NewInventory("", sims[tc.pool].CalculateLimit())
 		for i, swap := range tc.swaps {
 			t.Run(fmt.Sprintf("%v swap %d", tc.name, i), func(t *testing.T) {
-				res, err := sims[tc.pool].CalcAmountOut(pool.CalcAmountOutParams{
+				cloned := sims[tc.pool].CloneState()
+				calcAmountOutParams := pool.CalcAmountOutParams{
 					TokenAmountIn: pool.TokenAmount{
 						Token:  "A",
 						Amount: bignumber.NewBig10(swap.amountIn),
 					},
 					TokenOut: "B",
 					Limit:    limit,
-				})
+				}
+				res, err := sims[tc.pool].CalcAmountOut(calcAmountOutParams)
 
 				if swap.expOrderIds == nil {
 					require.NotNil(t, err)
@@ -1078,6 +1080,10 @@ func TestPool_UpdateBalance(t *testing.T) {
 
 				assert.Equal(t, swap.expAmountOut, res.TokenAmountOut.Amount.String())
 
+				clonedRes, err := cloned.CalcAmountOut(calcAmountOutParams)
+				require.Nil(t, err)
+				assert.Equal(t, res.TokenAmountOut, clonedRes.TokenAmountOut)
+
 				si := res.SwapInfo.(SwapInfo)
 				oid := make([]int64, 0, len(si.FilledOrders))
 				oinfo := ""
@@ -1101,6 +1107,10 @@ func TestPool_UpdateBalance(t *testing.T) {
 					SwapInfo:       res.SwapInfo,
 					SwapLimit:      limit,
 				})
+
+				clonedRes, err = cloned.CalcAmountOut(calcAmountOutParams)
+				require.Nil(t, err)
+				assert.Equal(t, res.TokenAmountOut, clonedRes.TokenAmountOut)
 			})
 		}
 	}

diff --git a/pkg/source/saddle/pool_simulator.go b/pkg/source/saddle/pool_simulator.go
@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/goccy/go-json"
+	"github.com/samber/lo"
 
 	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/entity"
 	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/source/pool"
@@ -193,6 +194,15 @@ func (t *PoolSimulator) CalcAmountOut(param pool.CalcAmountOutParams) (*pool.Cal
 	return &pool.CalcAmountOutResult{}, errors.New("i'm dead here")
 }
 
+func (t *PoolSimulator) CloneState() pool.IPoolSimulator {
+	cloned := *t
+	cloned.LpSupply = new(big.Int).Set(t.LpSupply)
+	cloned.Info.Reserves = lo.Map(t.Info.Reserves, func(v *big.Int, i int) *big.Int {
+		return new(big.Int).Set(v)
+	})
+	return &cloned
+}
+
 func (t *PoolSimulator) UpdateBalance(params pool.UpdateBalanceParams) {
 	input, output, fee := params.TokenAmountIn, params.TokenAmountOut, params.Fee
 	var inputAmount = input.Amount

diff --git a/pkg/source/saddle/pool_simulator_test.go b/pkg/source/saddle/pool_simulator_test.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/entity"
 	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/source/pool"
+	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/util/bignumber"
 	utils "github.com/KyberNetwork/kyberswap-dex-lib/pkg/util/bignumber"
 	"github.com/KyberNetwork/kyberswap-dex-lib/pkg/util/testutil"
 )
@@ -37,7 +38,8 @@ func TestCalcAmountOut_Saddle(t *testing.T) {
 	p, err := NewPoolSimulator(entity.Pool{
 		Exchange:    "",
 		Type:        "",
-		Reserves:    entity.PoolReserves{"64752405287155128155", "426593278742302082683", "66589357932477536907", "553429429583268691085"},
+		Reserves:    entity.PoolReserves{"64752405287155128155", "426593278742302082683", "66589357932477536907",
+			"553429429583268691085"},
 		Tokens:      []*entity.PoolToken{{Address: "A"}, {Address: "B"}, {Address: "C"}},
 		Extra:       "{\"initialA\":\"48000\",\"futureA\":\"92000\",\"initialATime\":1652287436,\"futureATime\":1653655053,\"swapFee\":\"4000000\",\"adminFee\":\"5000000000\"}",
 		StaticExtra: "{\"lpToken\":\"LP\",\"precisionMultipliers\":[\"1\",\"1\",\"1\"]}",
@@ -128,7 +130,8 @@ func TestCalcAmountOut_OneSwap(t *testing.T) {
 	p, err := NewPoolSimulator(entity.Pool{
 		Exchange:    "",
 		Type:        "",
-		Reserves:    entity.PoolReserves{"339028421564024338437", "347684462442560871352", "423798212946198474118", "315249216225911580289", "1404290718401538825321"},
+		Reserves:    entity.PoolReserves{"339028421564024338437", "347684462442560871352", "423798212946198474118",
+			"315249216225911580289", "1404290718401538825321"},
 		Tokens:      []*entity.PoolToken{{Address: "A"}, {Address: "B"}, {Address: "C"}, {Address: "D"}},
 		Extra:       "{\"initialA\":\"60000\",\"futureA\":\"60000\",\"initialATime\":0,\"futureATime\":0,\"swapFee\":\"1000000\",\"adminFee\":\"10000000000\",\"defaultWithdrawFee\":\"5000000\"}",
 		StaticExtra: "{\"lpToken\":\"LP\",\"precisionMultipliers\":[\"1\",\"1\",\"1\",\"1\"]}",
@@ -175,7 +178,8 @@ func TestCalcAmountOut_IronStable(t *testing.T) {
 	p, err := NewPoolSimulator(entity.Pool{
 		Exchange:    "",
 		Type:        "",
-		Reserves:    entity.PoolReserves{"233518765839", "198509040315", "228986742536043517345011", "654251953025609178732174"},
+		Reserves:    entity.PoolReserves{"233518765839", "198509040315", "228986742536043517345011",
+			"654251953025609178732174"},
 		Tokens:      []*entity.PoolToken{{Address: "A"}, {Address: "B"}, {Address: "C"}},
 		Extra:       "{\"initialA\":\"18000\",\"futureA\":\"120000\",\"initialATime\":1627094541,\"futureATime\":1627699238,\"swapFee\":\"2000000\",\"adminFee\":\"10000000000\", \"defaultWithdrawFee\":\"5000000\"}",
 		StaticExtra: "{\"lpToken\":\"LP\",\"precisionMultipliers\":[\"1000000000000\",\"1000000000000\",\"1\"]}",
@@ -207,13 +211,18 @@ func TestUpdateBalance_Saddle(t *testing.T) {
 	// test data from https://etherscan.io/address/0xa6018520eaacc06c30ff2e1b3ee2c7c22e64196a#readContract
 	testcases := []struct {
 		in               string
-		inAmount         int64
+		inAmount         string
 		out              string
 		expectedBalances []string
 	}{
-		{"A", 10000000, "B", []string{"64752405287165128155", "426593278742291992582", "66589357932477536907", "553429429583268691085"}},
-		{"A", 10000000, "C", []string{"64752405287175128155", "426593278742291992582", "66589357932467535940", "553429429583268691085"}},
-		{"B", 10000000, "A", []string{"64752405287165221417", "426593278742301992582", "66589357932467535940", "553429429583268691085"}},
+		{"A", "10000000", "B",
+			[]string{"64752405287165128155", "426593278742291992582", "66589357932477536907", "553429429583268691085"}},
+		{"A", "10000000", "C",
+			[]string{"64752405287175128155", "426593278742291992582", "66589357932467535940", "553429429583268691085"}},
+		{"B", "10000000", "A",
+			[]string{"64752405287165221417", "426593278742301992582", "66589357932467535940", "553429429583268691085"}},
+		{"C", "9500000000000000000", "B",
+			[]string{"64752405287165221417", "417021399572301197888", "76089357932467535940", "553429429583268691085"}},
 
 		// cannot test these case because we haven't accounted for token fee when adding/removing liq yet
 		// {"A", 10000000, "LP", []string{"64752405287175220754", "426593278742301992004", "66589357932467535849", "553429429583278677755"}},
@@ -222,7 +231,8 @@ func TestUpdateBalance_Saddle(t *testing.T) {
 	p, err := NewPoolSimulator(entity.Pool{
 		Exchange:    "",
 		Type:        "",
-		Reserves:    entity.PoolReserves{"64752405287155128155", "426593278742302082683", "66589357932477536907", "553429429583268691085"},
+		Reserves:    entity.PoolReserves{"64752405287155128155", "426593278742302082683", "66589357932477536907",
+			"553429429583268691085"},
 		Tokens:      []*entity.PoolToken{{Address: "A"}, {Address: "B"}, {Address: "C"}},
 		Extra:       "{\"initialA\":\"48000\",\"futureA\":\"92000\",\"initialATime\":1652287436,\"futureATime\":1653655053,\"swapFee\":\"4000000\",\"adminFee\":\"5000000000\"}",
 		StaticExtra: "{\"lpToken\":\"LP\",\"precisionMultipliers\":[\"1\",\"1\",\"1\"]}",
@@ -231,7 +241,8 @@ func TestUpdateBalance_Saddle(t *testing.T) {
 
 	for idx, tc := range testcases {
 		t.Run(fmt.Sprintf("test %d", idx), func(t *testing.T) {
-			amountIn := pool.TokenAmount{Token: tc.in, Amount: big.NewInt(tc.inAmount)}
+			cloned := p.CloneState()
+			amountIn := pool.TokenAmount{Token: tc.in, Amount: bignumber.NewBig10(tc.inAmount)}
 			out, err := testutil.MustConcurrentSafe[*pool.CalcAmountOutResult](t, func() (any, error) {
 				return p.CalcAmountOut(pool.CalcAmountOutParams{
 					TokenAmountIn: amountIn,
@@ -251,6 +262,13 @@ func TestUpdateBalance_Saddle(t *testing.T) {
 				assert.Equal(t, utils.NewBig10(tc.expectedBalances[i]), balance)
 			}
 			assert.Equal(t, utils.NewBig10(tc.expectedBalances[len(p.Info.Reserves)]), p.LpSupply)
+
+			clonedRes, err := cloned.CalcAmountOut(pool.CalcAmountOutParams{
+				TokenAmountIn: amountIn,
+				TokenOut:      tc.out,
+			})
+			require.Nil(t, err)
+			assert.Equal(t, clonedRes.TokenAmountOut, out.TokenAmountOut)
 		})
 	}
 }