wp-evm-amm-math 0.1.2

//! Tick <-> sqrt_price_x96 conversion.
//!
//! Faithful port of Uniswap V3 `TickMath.sol`.
//! All constants and magic numbers come directly from the
//! Solidity reference implementation.

use alloy_primitives::U256;

use crate::AmmMathError;

/// Minimum tick value.
pub const MIN_TICK: i32 = -887272;
/// Maximum tick value.
pub const MAX_TICK: i32 = 887272;

/// Minimum sqrt ratio (at MIN_TICK).
pub const MIN_SQRT_RATIO: U256 = U256::from_limbs([4295128739, 0, 0, 0]);

/// Maximum sqrt ratio (at MAX_TICK).
/// 1461446703485210103287273052203988822378723970342
pub const MAX_SQRT_RATIO: U256 =
    U256::from_limbs([6743328256752651558, 17280870778742802505, 4294805859, 0]);

/// Returns the sqrt ratio as a Q64.96 value for a given tick.
///
/// Equivalent to `TickMath.getSqrtRatioAtTick` in Solidity.
/// The result is a U256 representing sqrt(1.0001^tick) * 2^96.
pub fn get_sqrt_ratio_at_tick(tick: i32) -> crate::Result<U256> {
    if !(MIN_TICK..=MAX_TICK).contains(&tick) {
        return Err(AmmMathError::TickOutOfRange { tick });
    }

    let abs_tick = tick.unsigned_abs();

    let mut ratio: U256 = if abs_tick & 0x1 != 0 {
        U256::from_str_radix("fffcb933bd6fad37aa2d162d1a594001", 16).unwrap()
    } else {
        U256::from(1u64) << 128
    };

    macro_rules! apply_bit {
        ($bit:expr, $hex:expr) => {
            if abs_tick & (1 << $bit) != 0 {
                ratio = (ratio * U256::from_str_radix($hex, 16).unwrap()) >> 128;
            }
        };
    }

    apply_bit!(1, "fff97272373d413259a46990580e213a");
    apply_bit!(2, "fff2e50f5f656932ef12357cf3c7fdcc");
    apply_bit!(3, "ffe5caca7e10e4e61c3624eaa0941cd0");
    apply_bit!(4, "ffcb9843d60f6159c9db58835c926644");
    apply_bit!(5, "ff973b41fa98c081472e6896dfb254c0");
    apply_bit!(6, "ff2ea16466c96a3843ec78b326b52861");
    apply_bit!(7, "fe5dee046a99a2a811c461f1969c3053");
    apply_bit!(8, "fcbe86c7900a88aedcffc83b479aa3a4");
    apply_bit!(9, "f987a7253ac413176f2b074cf7815e54");
    apply_bit!(10, "f3392b0822b70005940c7a398e4b70f3");
    apply_bit!(11, "e7159475a2c29b7443b29c7fa6e889d9");
    apply_bit!(12, "d097f3bdfd2022b8845ad8f792aa5825");
    apply_bit!(13, "a9f746462d870fdf8a65dc1f90e061e5");
    apply_bit!(14, "70d869a156d2a1b890bb3df62baf32f7");
    apply_bit!(15, "31be135f97d08fd981231505542fcfa6");
    apply_bit!(16, "9aa508b5b7a84e1c677de54f3e99bc9");
    apply_bit!(17, "5d6af8dedb81196699c329225ee604");
    apply_bit!(18, "2216e584f5fa1ea926041bedfe98");
    apply_bit!(19, "48a170391f7dc42444e8fa2");

    // If tick is positive, invert the ratio.
    if tick > 0 {
        ratio = U256::MAX / ratio;
    }

    // Shift from Q128 down to Q96, rounding up if needed.
    let shift: U256 = ratio % (U256::from(1u64) << 32);
    let result = (ratio >> 32) + if shift.is_zero() { U256::ZERO } else { U256::from(1u64) };

    Ok(result)
}

/// Returns the tick corresponding to a given sqrt ratio.
///
/// Equivalent to `TickMath.getTickAtSqrtRatio` in Solidity.
/// The sqrt_price_x96 must be in [MIN_SQRT_RATIO, MAX_SQRT_RATIO).
///
/// Algorithm follows the optimized Aperture/uni-v3-lib approach
/// used by the Rust `uniswap-v3-sdk`.
pub fn get_tick_at_sqrt_ratio(sqrt_price_x96: U256) -> crate::Result<i32> {
    if sqrt_price_x96 < MIN_SQRT_RATIO || sqrt_price_x96 >= MAX_SQRT_RATIO {
        return Err(AmmMathError::SqrtPriceOutOfRange);
    }

    // Find MSB of sqrt_price_x96 (160 > msb >= 32).
    let msb = most_significant_bit(sqrt_price_x96);

    // Integer part of log_2 in 8.64 fixed point.
    // (msb - 96) << 64
    let mut log_2_x64: U256 = U256::from(msb as u64).wrapping_sub(U256::from(96u64)) << 64;

    // Normalize: r = sqrt_price_x96 << 96 >> (msb - 31)
    // This puts r in [2^127, 2^128).
    let mut r: U256 = (sqrt_price_x96 << 96) >> (msb - 31);

    // 14 iterations of binary log refinement.
    // Each iteration squares r (fits in U256 since r < 2^128),
    // checks if r^2 >= 2^255 (i.e. top bit set), and
    // accumulates the fractional bit.
    let mut decimals: u64 = 0;

    macro_rules! refine {
        ($bit_pos:expr, $last:expr) => {
            let square: U256 = r * r;
            let f = square.as_limbs()[3] >> 63;
            r = square >> (127 + f as u8);
            decimals |= f << $bit_pos;
        };
        ($bit_pos:expr) => {
            let square: U256 = r * r;
            let f = square.as_limbs()[3] >> 63;
            r = square >> (127 + f as u8);
            decimals |= f << $bit_pos;
        };
    }

    refine!(63);
    refine!(62);
    refine!(61);
    refine!(60);
    refine!(59);
    refine!(58);
    refine!(57);
    refine!(56);
    refine!(55);
    refine!(54);
    refine!(53);
    refine!(52);
    refine!(51);
    // Last iteration: don't need r anymore.
    {
        let square: U256 = r * r;
        let f = square.as_limbs()[3] >> 63;
        decimals |= f << 50;
    }

    log_2_x64 |= U256::from(decimals);

    // Convert log_2 to log_sqrt_10001.
    // tick = log_2 * (2^64 / log_2(sqrt(1.0001)))
    //      = log_2 * 255738958999603826347141
    let magic = U256::from_str_radix("255738958999603826347141", 10).unwrap();
    let log_sqrt10001: U256 = log_2_x64.wrapping_mul(magic);

    // tick_low and tick_high bracket the true tick.
    let sub_const = U256::from_str_radix("3402992956809132418596140100660247210", 10).unwrap();
    let add_const = U256::from_str_radix("291339464771989622907027621153398088495", 10).unwrap();

    let tick_low_u256: U256 = log_sqrt10001.wrapping_sub(sub_const) >> 128;
    let tick_low = tick_low_u256.as_limbs()[0] as i32;

    let tick_high_u256: U256 = log_sqrt10001.wrapping_add(add_const) >> 128;
    let tick_high = tick_high_u256.as_limbs()[0] as i32;

    if tick_low == tick_high {
        Ok(tick_low)
    } else {
        let sqrt_at_high = get_sqrt_ratio_at_tick(tick_high)?;
        if sqrt_at_high <= sqrt_price_x96 {
            Ok(tick_high)
        } else {
            Ok(tick_low)
        }
    }
}

/// Find the index of the most significant bit of a U256.
fn most_significant_bit(x: U256) -> u32 {
    255 - x.leading_zeros() as u32
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_tick_0() {
        let r = get_sqrt_ratio_at_tick(0).unwrap();
        assert_eq!(r, U256::from(1u64) << 96);
    }

    #[test]
    fn test_tick_1() {
        let r = get_sqrt_ratio_at_tick(1).unwrap();
        // Matches Solidity TickMath.sol and uniswap-v3-sdk:
        // ratio is rounded up when shifting Q128 -> Q96.
        assert_eq!(r, U256::from_str_radix("79232123823359799118286999568", 10).unwrap());
    }

    #[test]
    fn test_tick_neg_1() {
        let r = get_sqrt_ratio_at_tick(-1).unwrap();
        assert_eq!(r, U256::from_str_radix("79224201403219477170569942574", 10).unwrap());
    }

    #[test]
    fn test_min_tick() {
        let r = get_sqrt_ratio_at_tick(MIN_TICK).unwrap();
        assert_eq!(r, MIN_SQRT_RATIO);
    }

    #[test]
    fn test_max_tick() {
        let r = get_sqrt_ratio_at_tick(MAX_TICK).unwrap();
        assert_eq!(r, MAX_SQRT_RATIO);
    }

    #[test]
    fn test_out_of_range() {
        assert!(get_sqrt_ratio_at_tick(MIN_TICK - 1).is_err());
        assert!(get_sqrt_ratio_at_tick(MAX_TICK + 1).is_err());
    }

    #[test]
    fn test_get_tick_at_sqrt_ratio_tick_0() {
        let q96 = U256::from(1u64) << 96;
        let tick = get_tick_at_sqrt_ratio(q96).unwrap();
        assert_eq!(tick, 0);
    }

    #[test]
    fn test_get_tick_at_sqrt_ratio_tick_1() {
        // Use the actual output of get_sqrt_ratio_at_tick(1).
        let sqrt = get_sqrt_ratio_at_tick(1).unwrap();
        let tick = get_tick_at_sqrt_ratio(sqrt).unwrap();
        assert_eq!(tick, 1);
    }

    #[test]
    fn test_roundtrip_various_ticks() {
        // MAX_TICK is excluded because
        // get_sqrt_ratio_at_tick(MAX_TICK) == MAX_SQRT_RATIO
        // and get_tick_at_sqrt_ratio requires
        // sqrt_price < MAX_SQRT_RATIO.
        for t in [-887272, -50000, -1000, -100, -1, 0, 1, 100, 1000, 50000, 887271] {
            let sqrt = get_sqrt_ratio_at_tick(t).unwrap();
            let recovered = get_tick_at_sqrt_ratio(sqrt).unwrap();
            assert_eq!(recovered, t, "roundtrip failed for tick {t}");
        }
    }

    #[test]
    fn test_min_sqrt_ratio_value() {
        assert_eq!(MIN_SQRT_RATIO, U256::from(4295128739u64));
    }

    #[test]
    fn test_max_sqrt_ratio_value() {
        let expected =
            U256::from_str_radix("1461446703485210103287273052203988822378723970342", 10).unwrap();
        assert_eq!(MAX_SQRT_RATIO, expected);
    }

    #[test]
    fn test_roundtrip_large_ticks() {
        // Extend coverage to large tick values not in the
        // existing roundtrip test: -100000 and 100000.
        for t in [-100000i32, 100000] {
            let sqrt = get_sqrt_ratio_at_tick(t).unwrap();
            let recovered = get_tick_at_sqrt_ratio(sqrt).unwrap();
            assert_eq!(recovered, t, "roundtrip failed for tick {t}");
        }
    }

    #[test]
    fn fuzz_tick_to_sqrt_price_roundtrip() {
        use rand::{rngs::StdRng, SeedableRng};
        fn test_rng() -> StdRng {
            StdRng::from_os_rng()
        }
        use rand::Rng;

        let mut rng = test_rng();
        for _ in 0..1000 {
            // MAX_TICK is excluded: get_sqrt_ratio_at_tick(MAX_TICK) == MAX_SQRT_RATIO
            // and get_tick_at_sqrt_ratio requires sqrt_price < MAX_SQRT_RATIO.
            let tick: i32 = rng.random_range(MIN_TICK..MAX_TICK);
            let sqrt = get_sqrt_ratio_at_tick(tick).unwrap_or_else(|e| {
                panic!("get_sqrt_ratio_at_tick({tick}) failed: {e}");
            });
            let recovered = get_tick_at_sqrt_ratio(sqrt).unwrap_or_else(|e| {
                panic!("get_tick_at_sqrt_ratio failed for tick={tick}, sqrt={sqrt}: {e}");
            });
            assert_eq!(
                recovered, tick,
                "roundtrip failed: tick={tick}, sqrt={sqrt}, recovered={recovered}"
            );
        }
    }

    #[test]
    fn fuzz_sqrt_ratio_bounds() {
        use rand::{rngs::StdRng, SeedableRng};
        fn test_rng() -> StdRng {
            StdRng::from_os_rng()
        }
        use rand::Rng;

        let mut rng = test_rng();
        for _ in 0..1000 {
            let tick: i32 = rng.random_range(MIN_TICK..=MAX_TICK);
            let sqrt = get_sqrt_ratio_at_tick(tick).unwrap_or_else(|e| {
                panic!("get_sqrt_ratio_at_tick({tick}) failed: {e}");
            });
            assert!(
                sqrt >= MIN_SQRT_RATIO,
                "tick={tick} produced sqrt={sqrt} below MIN_SQRT_RATIO"
            );
            assert!(
                sqrt <= MAX_SQRT_RATIO,
                "tick={tick} produced sqrt={sqrt} above MAX_SQRT_RATIO"
            );
        }
    }

    #[test]
    fn fuzz_tick_monotonicity() {
        use rand::{rngs::StdRng, SeedableRng};
        fn test_rng() -> StdRng {
            StdRng::from_os_rng()
        }
        use rand::Rng;

        let mut rng = test_rng();
        for _ in 0..1000 {
            let tick_a: i32 = rng.random_range(MIN_TICK..MAX_TICK);
            let tick_b: i32 = rng.random_range((tick_a + 1)..=MAX_TICK);
            let sqrt_a = get_sqrt_ratio_at_tick(tick_a).unwrap();
            let sqrt_b = get_sqrt_ratio_at_tick(tick_b).unwrap();
            assert!(
                sqrt_a < sqrt_b,
                "monotonicity violated: tick_a={tick_a} sqrt_a={sqrt_a}, tick_b={tick_b} \
                 sqrt_b={sqrt_b}"
            );
        }
    }
}