mev-beta/pkg/parsers/uniswap_v3_math_test.go

package parsers

import (
	"math/big"
	"testing"
)

func TestGetSqrtRatioAtTick(t *testing.T) {
	tests := []struct {
		name    string
		tick    int32
		wantErr bool
	}{
		{
			name:    "tick 0 (price = 1)",
			tick:    0,
			wantErr: false,
		},
		{
			name:    "positive tick",
			tick:    100,
			wantErr: false,
		},
		{
			name:    "negative tick",
			tick:    -100,
			wantErr: false,
		},
		{
			name:    "max tick",
			tick:    MaxTick,
			wantErr: false,
		},
		{
			name:    "min tick",
			tick:    MinTick,
			wantErr: false,
		},
		{
			name:    "tick out of bounds (above)",
			tick:    MaxTick + 1,
			wantErr: true,
		},
		{
			name:    "tick out of bounds (below)",
			tick:    MinTick - 1,
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sqrtPrice, err := GetSqrtRatioAtTick(tt.tick)

			if tt.wantErr {
				if err == nil {
					t.Error("GetSqrtRatioAtTick() expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("GetSqrtRatioAtTick() unexpected error: %v", err)
			}

			if sqrtPrice == nil || sqrtPrice.Sign() <= 0 {
				t.Error("GetSqrtRatioAtTick() returned invalid sqrtPrice")
			}

			// Verify sqrtPrice is within valid range
			if sqrtPrice.Cmp(MinSqrtRatio) < 0 || sqrtPrice.Cmp(MaxSqrtRatio) > 0 {
				t.Errorf("GetSqrtRatioAtTick() sqrtPrice out of bounds: %v", sqrtPrice)
			}
		})
	}
}

func TestGetTickAtSqrtRatio(t *testing.T) {
	tests := []struct {
		name        string
		sqrtPriceX96 *big.Int
		wantErr     bool
	}{
		{
			name:        "Q96 (price = 1, tick ≈ 0)",
			sqrtPriceX96: Q96,
			wantErr:     false,
		},
		{
			name:        "min sqrt ratio",
			sqrtPriceX96: MinSqrtRatio,
			wantErr:     false,
		},
		{
			name:        "max sqrt ratio",
			sqrtPriceX96: MaxSqrtRatio,
			wantErr:     false,
		},
		{
			name:        "sqrt ratio below min",
			sqrtPriceX96: big.NewInt(1),
			wantErr:     true,
		},
		{
			name:        "sqrt ratio above max",
			sqrtPriceX96: new(big.Int).Add(MaxSqrtRatio, big.NewInt(1)),
			wantErr:     true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tick, err := GetTickAtSqrtRatio(tt.sqrtPriceX96)

			if tt.wantErr {
				if err == nil {
					t.Error("GetTickAtSqrtRatio() expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("GetTickAtSqrtRatio() unexpected error: %v", err)
			}

			// Verify tick is within valid range
			if tick < MinTick || tick > MaxTick {
				t.Errorf("GetTickAtSqrtRatio() tick out of bounds: %v", tick)
			}
		})
	}
}

func TestTickRoundTrip(t *testing.T) {
	// Test that tick -> sqrtPrice -> tick gives us back the same tick (or very close)
	testTicks := []int32{-100000, -10000, -1000, -100, 0, 100, 1000, 10000, 100000}

	for _, originalTick := range testTicks {
		t.Run("", func(t *testing.T) {
			sqrtPrice, err := GetSqrtRatioAtTick(originalTick)
			if err != nil {
				t.Fatalf("GetSqrtRatioAtTick() error: %v", err)
			}

			calculatedTick, err := GetTickAtSqrtRatio(sqrtPrice)
			if err != nil {
				t.Fatalf("GetTickAtSqrtRatio() error: %v", err)
			}

			// Allow for small rounding differences
			diff := originalTick - calculatedTick
			if diff < 0 {
				diff = -diff
			}

			if diff > 1 {
				t.Errorf("Tick round trip failed: original=%d, calculated=%d, diff=%d",
					originalTick, calculatedTick, diff)
			}
		})
	}
}

func TestGetAmount0Delta(t *testing.T) {
	tests := []struct {
		name        string
		sqrtRatioA  *big.Int
		sqrtRatioB  *big.Int
		liquidity   *big.Int
		roundUp     bool
		wantPositive bool
	}{
		{
			name:        "basic calculation",
			sqrtRatioA:  new(big.Int).Lsh(big.NewInt(1), 96), // Q96
			sqrtRatioB:  new(big.Int).Lsh(big.NewInt(2), 96), // 2 * Q96
			liquidity:   big.NewInt(1000000),
			roundUp:     false,
			wantPositive: true,
		},
		{
			name:        "same ratios (zero delta)",
			sqrtRatioA:  Q96,
			sqrtRatioB:  Q96,
			liquidity:   big.NewInt(1000000),
			roundUp:     false,
			wantPositive: false,
		},
		{
			name:        "zero liquidity",
			sqrtRatioA:  Q96,
			sqrtRatioB:  new(big.Int).Lsh(big.NewInt(2), 96),
			liquidity:   big.NewInt(0),
			roundUp:     false,
			wantPositive: false,
		},
		{
			name:        "round up",
			sqrtRatioA:  Q96,
			sqrtRatioB:  new(big.Int).Lsh(big.NewInt(2), 96),
			liquidity:   big.NewInt(1000000),
			roundUp:     true,
			wantPositive: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			amount := GetAmount0Delta(tt.sqrtRatioA, tt.sqrtRatioB, tt.liquidity, tt.roundUp)

			if tt.wantPositive && amount.Sign() <= 0 {
				t.Error("GetAmount0Delta() expected positive amount, got zero or negative")
			}

			if !tt.wantPositive && amount.Sign() > 0 {
				t.Error("GetAmount0Delta() expected zero or negative amount, got positive")
			}
		})
	}
}

func TestGetAmount1Delta(t *testing.T) {
	tests := []struct {
		name        string
		sqrtRatioA  *big.Int
		sqrtRatioB  *big.Int
		liquidity   *big.Int
		roundUp     bool
		wantPositive bool
	}{
		{
			name:        "basic calculation",
			sqrtRatioA:  Q96,
			sqrtRatioB:  new(big.Int).Lsh(big.NewInt(2), 96),
			liquidity:   big.NewInt(1000000),
			roundUp:     false,
			wantPositive: true,
		},
		{
			name:        "same ratios (zero delta)",
			sqrtRatioA:  Q96,
			sqrtRatioB:  Q96,
			liquidity:   big.NewInt(1000000),
			roundUp:     false,
			wantPositive: false,
		},
		{
			name:        "zero liquidity",
			sqrtRatioA:  Q96,
			sqrtRatioB:  new(big.Int).Lsh(big.NewInt(2), 96),
			liquidity:   big.NewInt(0),
			roundUp:     false,
			wantPositive: false,
		},
		{
			name:        "round up",
			sqrtRatioA:  Q96,
			sqrtRatioB:  new(big.Int).Lsh(big.NewInt(2), 96),
			liquidity:   big.NewInt(1000000),
			roundUp:     true,
			wantPositive: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			amount := GetAmount1Delta(tt.sqrtRatioA, tt.sqrtRatioB, tt.liquidity, tt.roundUp)

			if tt.wantPositive && amount.Sign() <= 0 {
				t.Error("GetAmount1Delta() expected positive amount, got zero or negative")
			}

			if !tt.wantPositive && amount.Sign() > 0 {
				t.Error("GetAmount1Delta() expected zero or negative amount, got positive")
			}
		})
	}
}

func TestGetNextSqrtPriceFromInput(t *testing.T) {
	tests := []struct {
		name       string
		sqrtPrice  *big.Int
		liquidity  *big.Int
		amountIn   *big.Int
		zeroForOne bool
		wantErr    bool
	}{
		{
			name:       "swap token0 for token1",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000),
			amountIn:   big.NewInt(1000),
			zeroForOne: true,
			wantErr:    false,
		},
		{
			name:       "swap token1 for token0",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000),
			amountIn:   big.NewInt(1000),
			zeroForOne: false,
			wantErr:    false,
		},
		{
			name:       "zero liquidity",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(0),
			amountIn:   big.NewInt(1000),
			zeroForOne: true,
			wantErr:    true,
		},
		{
			name:       "zero sqrt price",
			sqrtPrice:  big.NewInt(0),
			liquidity:  big.NewInt(1000000),
			amountIn:   big.NewInt(1000),
			zeroForOne: true,
			wantErr:    true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			nextPrice, err := GetNextSqrtPriceFromInput(tt.sqrtPrice, tt.liquidity, tt.amountIn, tt.zeroForOne)

			if tt.wantErr {
				if err == nil {
					t.Error("GetNextSqrtPriceFromInput() expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("GetNextSqrtPriceFromInput() unexpected error: %v", err)
			}

			if nextPrice == nil || nextPrice.Sign() <= 0 {
				t.Error("GetNextSqrtPriceFromInput() returned invalid price")
			}

			// Verify price changed
			if nextPrice.Cmp(tt.sqrtPrice) == 0 {
				t.Error("GetNextSqrtPriceFromInput() price did not change")
			}

			// Verify price moved in correct direction
			if tt.zeroForOne {
				// Swapping token0 for token1 should decrease price
				if nextPrice.Cmp(tt.sqrtPrice) >= 0 {
					t.Error("GetNextSqrtPriceFromInput() price should decrease for zeroForOne swap")
				}
			} else {
				// Swapping token1 for token0 should increase price
				if nextPrice.Cmp(tt.sqrtPrice) <= 0 {
					t.Error("GetNextSqrtPriceFromInput() price should increase for oneForZero swap")
				}
			}
		})
	}
}

func TestGetNextSqrtPriceFromOutput(t *testing.T) {
	tests := []struct {
		name       string
		sqrtPrice  *big.Int
		liquidity  *big.Int
		amountOut  *big.Int
		zeroForOne bool
		wantErr    bool
	}{
		{
			name:       "swap token0 for token1 (output token1)",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000),
			amountOut:  big.NewInt(100),
			zeroForOne: true,
			wantErr:    false,
		},
		{
			name:       "swap token1 for token0 (output token0)",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000),
			amountOut:  big.NewInt(100),
			zeroForOne: false,
			wantErr:    false,
		},
		{
			name:       "zero liquidity",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(0),
			amountOut:  big.NewInt(100),
			zeroForOne: true,
			wantErr:    true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			nextPrice, err := GetNextSqrtPriceFromOutput(tt.sqrtPrice, tt.liquidity, tt.amountOut, tt.zeroForOne)

			if tt.wantErr {
				if err == nil {
					t.Error("GetNextSqrtPriceFromOutput() expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("GetNextSqrtPriceFromOutput() unexpected error: %v", err)
			}

			if nextPrice == nil || nextPrice.Sign() <= 0 {
				t.Error("GetNextSqrtPriceFromOutput() returned invalid price")
			}
		})
	}
}

func TestComputeSwapStep(t *testing.T) {
	sqrtPriceCurrent := Q96 // Price = 1
	sqrtPriceTarget := new(big.Int).Lsh(big.NewInt(2), 96) // Price = 2
	liquidity := big.NewInt(1000000000000) // 1 trillion
	amountRemaining := big.NewInt(1000000000000000000) // 1 ETH
	feePips := uint32(3000) // 0.3%

	sqrtPriceNext, amountIn, amountOut, feeAmount, err := ComputeSwapStep(
		sqrtPriceCurrent,
		sqrtPriceTarget,
		liquidity,
		amountRemaining,
		feePips,
	)

	if err != nil {
		t.Fatalf("ComputeSwapStep() error: %v", err)
	}

	if sqrtPriceNext == nil || sqrtPriceNext.Sign() <= 0 {
		t.Error("ComputeSwapStep() returned invalid sqrtPriceNext")
	}

	if amountIn == nil || amountIn.Sign() < 0 {
		t.Error("ComputeSwapStep() returned invalid amountIn")
	}

	if amountOut == nil || amountOut.Sign() <= 0 {
		t.Error("ComputeSwapStep() returned invalid amountOut")
	}

	if feeAmount == nil || feeAmount.Sign() < 0 {
		t.Error("ComputeSwapStep() returned invalid feeAmount")
	}

	t.Logf("Swap step results:")
	t.Logf("  sqrtPriceNext: %v", sqrtPriceNext)
	t.Logf("  amountIn: %v", amountIn)
	t.Logf("  amountOut: %v", amountOut)
	t.Logf("  feeAmount: %v", feeAmount)
}

func TestCalculateSwapAmounts(t *testing.T) {
	tests := []struct {
		name       string
		sqrtPrice  *big.Int
		liquidity  *big.Int
		amountIn   *big.Int
		zeroForOne bool
		feePips    uint32
		wantErr    bool
	}{
		{
			name:       "swap 1 token0 for token1",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000000000),
			amountIn:   big.NewInt(1000000000000000000), // 1 ETH
			zeroForOne: true,
			feePips:    3000, // 0.3%
			wantErr:    false,
		},
		{
			name:       "swap 1 token1 for token0",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000000000),
			amountIn:   big.NewInt(1000000), // 1 USDC (6 decimals)
			zeroForOne: false,
			feePips:    3000,
			wantErr:    false,
		},
		{
			name:       "high fee tier (1%)",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000000000),
			amountIn:   big.NewInt(1000000000000000000),
			zeroForOne: true,
			feePips:    10000, // 1%
			wantErr:    false,
		},
		{
			name:       "low fee tier (0.05%)",
			sqrtPrice:  Q96,
			liquidity:  big.NewInt(1000000000000),
			amountIn:   big.NewInt(1000000000000000000),
			zeroForOne: true,
			feePips:    500, // 0.05%
			wantErr:    false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			amountOut, priceAfter, err := CalculateSwapAmounts(
				tt.sqrtPrice,
				tt.liquidity,
				tt.amountIn,
				tt.zeroForOne,
				tt.feePips,
			)

			if tt.wantErr {
				if err == nil {
					t.Error("CalculateSwapAmounts() expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("CalculateSwapAmounts() unexpected error: %v", err)
			}

			if amountOut == nil || amountOut.Sign() <= 0 {
				t.Error("CalculateSwapAmounts() returned invalid amountOut")
			}

			if priceAfter == nil || priceAfter.Sign() <= 0 {
				t.Error("CalculateSwapAmounts() returned invalid priceAfter")
			}

			// Verify price moved in correct direction
			if tt.zeroForOne {
				if priceAfter.Cmp(tt.sqrtPrice) >= 0 {
					t.Error("CalculateSwapAmounts() price should decrease for zeroForOne swap")
				}
			} else {
				if priceAfter.Cmp(tt.sqrtPrice) <= 0 {
					t.Error("CalculateSwapAmounts() price should increase for oneForZero swap")
				}
			}

			t.Logf("Swap results:")
			t.Logf("  amountIn: %v", tt.amountIn)
			t.Logf("  amountOut: %v", amountOut)
			t.Logf("  priceBefore: %v", tt.sqrtPrice)
			t.Logf("  priceAfter: %v", priceAfter)
		})
	}
}

func TestKnownPoolState(t *testing.T) {
	// Test with known values from a real Uniswap V3 pool
	// Example: WETH/USDC 0.3% pool on Arbitrum

	// At tick 0, price = 1
	tick := int32(0)
	sqrtPrice, err := GetSqrtRatioAtTick(tick)
	if err != nil {
		t.Fatalf("GetSqrtRatioAtTick() error: %v", err)
	}

	// SqrtPrice at tick 0 should be approximately Q96
	expectedSqrtPrice := Q96
	tolerance := new(big.Int).Div(Q96, big.NewInt(100)) // 1% tolerance

	diff := new(big.Int).Sub(sqrtPrice, expectedSqrtPrice)
	if diff.Sign() < 0 {
		diff.Neg(diff)
	}

	if diff.Cmp(tolerance) > 0 {
		t.Errorf("SqrtPrice at tick 0 not close to Q96: got %v, want %v, diff %v",
			sqrtPrice, expectedSqrtPrice, diff)
	}

	// Reverse calculation should give us back tick 0
	calculatedTick, err := GetTickAtSqrtRatio(sqrtPrice)
	if err != nil {
		t.Fatalf("GetTickAtSqrtRatio() error: %v", err)
	}

	if calculatedTick != tick && calculatedTick != tick-1 && calculatedTick != tick+1 {
		t.Errorf("Tick round trip failed: original=%d, calculated=%d", tick, calculatedTick)
	}
}