mev-beta/pkg/parsers/uniswap_v3_math.go

package parsers

import (
	"errors"
	"math"
	"math/big"
)

// Uniswap V3 Math Utilities
// Based on: https://github.com/Uniswap/v3-core and https://github.com/t4sk/clamm
//
// Key Constants:
// - Q96 = 2^96 (fixed-point precision for sqrtPriceX96)
// - MIN_TICK = -887272
// - MAX_TICK = 887272
// - MIN_SQRT_RATIO = 4295128739
// - MAX_SQRT_RATIO = 1461446703485210103287273052203988822378723970342

var (
	// Q96 is 2^96 for fixed-point arithmetic
	Q96 = new(big.Int).Lsh(big.NewInt(1), 96)

	// Q128 is 2^128
	Q128 = new(big.Int).Lsh(big.NewInt(1), 128)

	// Tick bounds
	MinTick int32 = -887272
	MaxTick int32 = 887272

	// SqrtPrice bounds (Q64.96 format)
	MinSqrtRatio = big.NewInt(4295128739)
	MaxSqrtRatio = mustParseBigInt("1461446703485210103287273052203988822378723970342")

	// 1.0001 as a ratio for tick calculations
	// TickBase = 1.0001 (the ratio between adjacent ticks)
	TickBase = 1.0001

	// Error definitions
	ErrInvalidTick          = errors.New("tick out of bounds")
	ErrInvalidSqrtPrice     = errors.New("sqrt price out of bounds")
	ErrInvalidLiquidity     = errors.New("liquidity must be positive")
	ErrPriceLimitReached    = errors.New("price limit reached")
)

// mustParseBigInt parses a decimal string to big.Int, panics on error
func mustParseBigInt(s string) *big.Int {
	n := new(big.Int)
	n.SetString(s, 10)
	return n
}

// GetSqrtRatioAtTick calculates sqrtPriceX96 from tick
// Formula: sqrt(1.0001^tick) * 2^96
func GetSqrtRatioAtTick(tick int32) (*big.Int, error) {
	if tick < MinTick || tick > MaxTick {
		return nil, ErrInvalidTick
	}

	// Calculate 1.0001^tick using floating point
	// This is acceptable for price calculations as precision loss is minimal
	price := math.Pow(TickBase, float64(tick))
	sqrtPrice := math.Sqrt(price)

	// Convert to Q96 format
	sqrtPriceX96Float := sqrtPrice * math.Pow(2, 96)

	// Convert to big.Int
	sqrtPriceX96 := new(big.Float).SetFloat64(sqrtPriceX96Float)
	result, _ := sqrtPriceX96.Int(nil)

	return result, nil
}

// GetTickAtSqrtRatio calculates tick from sqrtPriceX96
// Formula: tick = floor(log_1.0001(price)) = floor(log(price) / log(1.0001))
func GetTickAtSqrtRatio(sqrtPriceX96 *big.Int) (int32, error) {
	if sqrtPriceX96.Cmp(MinSqrtRatio) < 0 || sqrtPriceX96.Cmp(MaxSqrtRatio) > 0 {
		return 0, ErrInvalidSqrtPrice
	}

	// Convert Q96 to float
	sqrtPriceFloat := new(big.Float).SetInt(sqrtPriceX96)
	q96Float := new(big.Float).SetInt(Q96)
	sqrtPrice := new(big.Float).Quo(sqrtPriceFloat, q96Float)

	sqrtPriceF64, _ := sqrtPrice.Float64()
	price := sqrtPriceF64 * sqrtPriceF64

	// Calculate tick = log(price) / log(1.0001)
	tick := math.Log(price) / math.Log(TickBase)

	return int32(math.Floor(tick)), nil
}

// GetAmount0Delta calculates the amount0 delta for a liquidity change
// Formula: amount0 = liquidity * (sqrtRatioB - sqrtRatioA) / (sqrtRatioA * sqrtRatioB)
// When liquidity increases (adding), amount0 is positive
// When liquidity decreases (removing), amount0 is negative
func GetAmount0Delta(sqrtRatioA, sqrtRatioB, liquidity *big.Int, roundUp bool) *big.Int {
	if sqrtRatioA.Cmp(sqrtRatioB) > 0 {
		sqrtRatioA, sqrtRatioB = sqrtRatioB, sqrtRatioA
	}

	if liquidity.Sign() <= 0 {
		return big.NewInt(0)
	}

	// numerator = liquidity * (sqrtRatioB - sqrtRatioA) * 2^96
	numerator := new(big.Int).Sub(sqrtRatioB, sqrtRatioA)
	numerator.Mul(numerator, liquidity)
	numerator.Lsh(numerator, 96)

	// denominator = sqrtRatioA * sqrtRatioB
	denominator := new(big.Int).Mul(sqrtRatioA, sqrtRatioB)

	if roundUp {
		// Round up: (numerator + denominator - 1) / denominator
		result := new(big.Int).Sub(denominator, big.NewInt(1))
		result.Add(result, numerator)
		result.Div(result, denominator)
		return result
	}

	// Round down: numerator / denominator
	return new(big.Int).Div(numerator, denominator)
}

// GetAmount1Delta calculates the amount1 delta for a liquidity change
// Formula: amount1 = liquidity * (sqrtRatioB - sqrtRatioA) / 2^96
// When liquidity increases (adding), amount1 is positive
// When liquidity decreases (removing), amount1 is negative
func GetAmount1Delta(sqrtRatioA, sqrtRatioB, liquidity *big.Int, roundUp bool) *big.Int {
	if sqrtRatioA.Cmp(sqrtRatioB) > 0 {
		sqrtRatioA, sqrtRatioB = sqrtRatioB, sqrtRatioA
	}

	if liquidity.Sign() <= 0 {
		return big.NewInt(0)
	}

	// amount1 = liquidity * (sqrtRatioB - sqrtRatioA) / 2^96
	diff := new(big.Int).Sub(sqrtRatioB, sqrtRatioA)
	result := new(big.Int).Mul(liquidity, diff)

	if roundUp {
		// Round up: (result + Q96 - 1) / Q96
		result.Add(result, new(big.Int).Sub(Q96, big.NewInt(1)))
	}

	result.Rsh(result, 96)
	return result
}

// GetNextSqrtPriceFromInput calculates the next sqrtPrice given an input amount
// Used for exact input swaps
// zeroForOne: true if swapping token0 for token1, false otherwise
func GetNextSqrtPriceFromInput(sqrtPriceX96, liquidity, amountIn *big.Int, zeroForOne bool) (*big.Int, error) {
	if sqrtPriceX96.Sign() <= 0 || liquidity.Sign() <= 0 {
		return nil, ErrInvalidLiquidity
	}

	if zeroForOne {
		// Swapping token0 for token1
		// sqrtP' = (liquidity * sqrtP) / (liquidity + amountIn * sqrtP / 2^96)
		return getNextSqrtPriceFromAmount0RoundingUp(sqrtPriceX96, liquidity, amountIn, true)
	}

	// Swapping token1 for token0
	// sqrtP' = sqrtP + (amountIn * 2^96) / liquidity
	return getNextSqrtPriceFromAmount1RoundingDown(sqrtPriceX96, liquidity, amountIn, true)
}

// GetNextSqrtPriceFromOutput calculates the next sqrtPrice given an output amount
// Used for exact output swaps
// zeroForOne: true if swapping token0 for token1, false otherwise
func GetNextSqrtPriceFromOutput(sqrtPriceX96, liquidity, amountOut *big.Int, zeroForOne bool) (*big.Int, error) {
	if sqrtPriceX96.Sign() <= 0 || liquidity.Sign() <= 0 {
		return nil, ErrInvalidLiquidity
	}

	if zeroForOne {
		// Swapping token0 for token1 (outputting token1)
		// sqrtP' = sqrtP - (amountOut * 2^96) / liquidity
		return getNextSqrtPriceFromAmount1RoundingDown(sqrtPriceX96, liquidity, amountOut, false)
	}

	// Swapping token1 for token0 (outputting token0)
	// sqrtP' = (liquidity * sqrtP) / (liquidity - amountOut * sqrtP / 2^96)
	return getNextSqrtPriceFromAmount0RoundingUp(sqrtPriceX96, liquidity, amountOut, false)
}

// getNextSqrtPriceFromAmount0RoundingUp helper for amount0 calculations
func getNextSqrtPriceFromAmount0RoundingUp(sqrtPriceX96, liquidity, amount *big.Int, add bool) (*big.Int, error) {
	if amount.Sign() == 0 {
		return sqrtPriceX96, nil
	}

	// numerator = liquidity * sqrtPriceX96 * 2^96
	numerator := new(big.Int).Mul(liquidity, sqrtPriceX96)
	numerator.Lsh(numerator, 96)

	// product = amount * sqrtPriceX96
	product := new(big.Int).Mul(amount, sqrtPriceX96)

	if add {
		// denominator = liquidity * 2^96 + product
		denominator := new(big.Int).Lsh(liquidity, 96)
		denominator.Add(denominator, product)

		// Check for overflow
		if denominator.Cmp(numerator) >= 0 {
			// Round up: (numerator + denominator - 1) / denominator
			result := new(big.Int).Sub(denominator, big.NewInt(1))
			result.Add(result, numerator)
			result.Div(result, denominator)
			return result, nil
		}
	} else {
		// denominator = liquidity * 2^96 - product
		denominator := new(big.Int).Lsh(liquidity, 96)
		if product.Cmp(denominator) >= 0 {
			return nil, ErrPriceLimitReached
		}
		denominator.Sub(denominator, product)

		// Round up: (numerator + denominator - 1) / denominator
		result := new(big.Int).Sub(denominator, big.NewInt(1))
		result.Add(result, numerator)
		result.Div(result, denominator)
		return result, nil
	}

	// Fallback calculation
	return new(big.Int).Div(numerator, new(big.Int).Lsh(liquidity, 96)), nil
}

// getNextSqrtPriceFromAmount1RoundingDown helper for amount1 calculations
func getNextSqrtPriceFromAmount1RoundingDown(sqrtPriceX96, liquidity, amount *big.Int, add bool) (*big.Int, error) {
	if add {
		// sqrtP' = sqrtP + (amount * 2^96) / liquidity
		quotient := new(big.Int).Lsh(amount, 96)
		quotient.Div(quotient, liquidity)

		result := new(big.Int).Add(sqrtPriceX96, quotient)
		return result, nil
	}

	// sqrtP' = sqrtP - (amount * 2^96) / liquidity
	quotient := new(big.Int).Lsh(amount, 96)
	quotient.Div(quotient, liquidity)

	if quotient.Cmp(sqrtPriceX96) >= 0 {
		return nil, ErrPriceLimitReached
	}

	result := new(big.Int).Sub(sqrtPriceX96, quotient)
	return result, nil
}

// ComputeSwapStep simulates a single swap step within a tick range
// Returns: sqrtPriceX96Next, amountIn, amountOut, feeAmount
func ComputeSwapStep(
	sqrtRatioCurrentX96 *big.Int,
	sqrtRatioTargetX96 *big.Int,
	liquidity *big.Int,
	amountRemaining *big.Int,
	feePips uint32, // Fee in pips (1/1000000), e.g., 3000 = 0.3%
) (*big.Int, *big.Int, *big.Int, *big.Int, error) {
	zeroForOne := sqrtRatioCurrentX96.Cmp(sqrtRatioTargetX96) >= 0
	exactIn := amountRemaining.Sign() >= 0

	var sqrtRatioNextX96 *big.Int
	var amountIn, amountOut, feeAmount *big.Int

	if exactIn {
		// Calculate fee
		amountRemainingLessFee := new(big.Int).Mul(
			amountRemaining,
			big.NewInt(int64(1000000-feePips)),
		)
		amountRemainingLessFee.Div(amountRemainingLessFee, big.NewInt(1000000))

		// Calculate max amount we can swap in this step
		if zeroForOne {
			amountIn = GetAmount0Delta(sqrtRatioTargetX96, sqrtRatioCurrentX96, liquidity, true)
		} else {
			amountIn = GetAmount1Delta(sqrtRatioCurrentX96, sqrtRatioTargetX96, liquidity, true)
		}

		// Determine if we can complete the swap in this step
		if amountRemainingLessFee.Cmp(amountIn) >= 0 {
			// We can complete the swap, use target price
			sqrtRatioNextX96 = sqrtRatioTargetX96
		} else {
			// We cannot complete the swap, calculate new price
			var err error
			sqrtRatioNextX96, err = GetNextSqrtPriceFromInput(
				sqrtRatioCurrentX96,
				liquidity,
				amountRemainingLessFee,
				zeroForOne,
			)
			if err != nil {
				return nil, nil, nil, nil, err
			}
		}

		// Calculate amounts
		if zeroForOne {
			amountIn = GetAmount0Delta(sqrtRatioNextX96, sqrtRatioCurrentX96, liquidity, true)
			amountOut = GetAmount1Delta(sqrtRatioNextX96, sqrtRatioCurrentX96, liquidity, false)
		} else {
			amountIn = GetAmount1Delta(sqrtRatioCurrentX96, sqrtRatioNextX96, liquidity, true)
			amountOut = GetAmount0Delta(sqrtRatioNextX96, sqrtRatioCurrentX96, liquidity, false)
		}

		// Calculate fee
		if sqrtRatioNextX96.Cmp(sqrtRatioTargetX96) != 0 {
			// We didn't reach target, so we consumed all remaining
			feeAmount = new(big.Int).Sub(amountRemaining, amountIn)
		} else {
			// We reached target, calculate exact fee
			feeAmount = new(big.Int).Mul(amountIn, big.NewInt(int64(feePips)))
			feeAmount.Div(feeAmount, big.NewInt(int64(1000000-feePips)))
			// Round up
			feeAmount.Add(feeAmount, big.NewInt(1))
		}
	} else {
		// Exact output swap (not commonly used in MEV)
		// Implementation simplified for now
		sqrtRatioNextX96 = sqrtRatioTargetX96
		amountIn = big.NewInt(0)
		amountOut = new(big.Int).Abs(amountRemaining)
		feeAmount = big.NewInt(0)
	}

	return sqrtRatioNextX96, amountIn, amountOut, feeAmount, nil
}

// CalculateSwapAmounts calculates the output amount for a given input amount
// This is useful for simulating swaps and calculating expected profits
func CalculateSwapAmounts(
	sqrtPriceX96 *big.Int,
	liquidity *big.Int,
	amountIn *big.Int,
	zeroForOne bool,
	feePips uint32,
) (amountOut *big.Int, priceAfter *big.Int, err error) {
	// Subtract fee from input
	amountInAfterFee := new(big.Int).Mul(amountIn, big.NewInt(int64(1000000-feePips)))
	amountInAfterFee.Div(amountInAfterFee, big.NewInt(1000000))

	// Calculate new sqrt price
	priceAfter, err = GetNextSqrtPriceFromInput(sqrtPriceX96, liquidity, amountInAfterFee, zeroForOne)
	if err != nil {
		return nil, nil, err
	}

	// Calculate output amount
	if zeroForOne {
		amountOut = GetAmount1Delta(priceAfter, sqrtPriceX96, liquidity, false)
	} else {
		amountOut = GetAmount0Delta(priceAfter, sqrtPriceX96, liquidity, false)
	}

	// Ensure output is positive
	if amountOut.Sign() < 0 {
		amountOut.Neg(amountOut)
	}

	return amountOut, priceAfter, nil
}