wp-evm-amm-math 0.1.2

//! Liquidity math: max liquidity for given token amounts.
//!
//! Port of Uniswap V3 `LiquidityAmounts.sol` (periphery).

use alloy_primitives::U256;

use crate::{full_math::mul_div, AmmMathError};

/// Q96 constant (2^96).
const Q96: U256 = U256::from_limbs([0, 1 << 32, 0, 0]);

/// Compute the liquidity amount for a given amount of token0
/// and a price range [sqrt_a, sqrt_b].
///
/// `L = amount0 * sqrt_a * sqrt_b / (sqrt_b - sqrt_a)`
pub fn liquidity_for_amount_0(
    sqrt_ratio_a_x96: U256,
    sqrt_ratio_b_x96: U256,
    amount0: U256,
) -> crate::Result<u128> {
    let (sqrt_lower, sqrt_upper) = if sqrt_ratio_a_x96 <= sqrt_ratio_b_x96 {
        (sqrt_ratio_a_x96, sqrt_ratio_b_x96)
    } else {
        (sqrt_ratio_b_x96, sqrt_ratio_a_x96)
    };

    let diff = sqrt_upper - sqrt_lower;
    if diff.is_zero() {
        return Ok(0);
    }

    // intermediate = amount0 * sqrt_lower * sqrt_upper / Q96
    let intermediate = mul_div(amount0, sqrt_lower, Q96)?;
    let liquidity = mul_div(intermediate, sqrt_upper, diff)?;

    if liquidity > U256::from(u128::MAX) {
        Err(AmmMathError::LiquidityOverflow)
    } else {
        Ok(u128::from(liquidity.as_limbs()[0]))
    }
}

/// Compute the liquidity amount for a given amount of token1
/// and a price range [sqrt_a, sqrt_b].
///
/// `L = amount1 / (sqrt_b - sqrt_a) * Q96`
pub fn liquidity_for_amount_1(
    sqrt_ratio_a_x96: U256,
    sqrt_ratio_b_x96: U256,
    amount1: U256,
) -> crate::Result<u128> {
    let (sqrt_lower, sqrt_upper) = if sqrt_ratio_a_x96 <= sqrt_ratio_b_x96 {
        (sqrt_ratio_a_x96, sqrt_ratio_b_x96)
    } else {
        (sqrt_ratio_b_x96, sqrt_ratio_a_x96)
    };

    let diff = sqrt_upper - sqrt_lower;
    if diff.is_zero() {
        return Ok(0);
    }

    let liquidity = mul_div(amount1, Q96, diff)?;

    if liquidity > U256::from(u128::MAX) {
        return Err(AmmMathError::LiquidityOverflow);
    }

    Ok(liquidity.to::<u128>())
}

/// Compute the maximum liquidity that can be provided given
/// amounts of both tokens at a current price within a range.
///
/// This is the periphery function
/// `LiquidityAmounts.getLiquidityForAmounts`.
pub fn max_liquidity_for_amounts(
    sqrt_ratio_x96: U256,
    sqrt_ratio_a_x96: U256,
    sqrt_ratio_b_x96: U256,
    amount0: U256,
    amount1: U256,
) -> crate::Result<u128> {
    let (sqrt_lower, sqrt_upper) = if sqrt_ratio_a_x96 <= sqrt_ratio_b_x96 {
        (sqrt_ratio_a_x96, sqrt_ratio_b_x96)
    } else {
        (sqrt_ratio_b_x96, sqrt_ratio_a_x96)
    };

    if sqrt_ratio_x96 <= sqrt_lower {
        // Price below range: only token0 matters.
        liquidity_for_amount_0(sqrt_lower, sqrt_upper, amount0)
    } else if sqrt_ratio_x96 < sqrt_upper {
        // Price inside range: take the minimum.
        let l0 = liquidity_for_amount_0(sqrt_ratio_x96, sqrt_upper, amount0)?;
        let l1 = liquidity_for_amount_1(sqrt_lower, sqrt_ratio_x96, amount1)?;
        Ok(l0.min(l1))
    } else {
        // Price above range: only token1 matters.
        liquidity_for_amount_1(sqrt_lower, sqrt_upper, amount1)
    }
}

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

    #[test]
    fn ref_vector_max_liquidity_inside_range() {
        // Oracle: sp=Q96, lo=Q96-1e27, hi=Q96+1e27, amount0=amount1=1e9.
        // Price inside range → min(l0, l1) = 79228162514.
        let sp = U256::from_str_radix("79228162514264337593543950336", 10).unwrap();
        let lo = U256::from_str_radix("78228162514264337593543950336", 10).unwrap();
        let hi = U256::from_str_radix("80228162514264337593543950336", 10).unwrap();
        let l = max_liquidity_for_amounts(
            sp,
            lo,
            hi,
            U256::from(1_000_000_000_u64),
            U256::from(1_000_000_000_u64),
        )
        .unwrap();
        assert_eq!(l, 79_228_162_514_u128);
    }

    #[test]
    fn test_liquidity_for_amount_0_basic() {
        let sqrt_a = get_sqrt_ratio_at_tick(0).unwrap();
        let sqrt_b = get_sqrt_ratio_at_tick(100).unwrap();
        let amount0 = U256::from(1_000_000u64);
        let liq = liquidity_for_amount_0(sqrt_a, sqrt_b, amount0).unwrap();
        assert!(liq > 0);
    }

    #[test]
    fn test_liquidity_for_amount_1_basic() {
        let sqrt_a = get_sqrt_ratio_at_tick(0).unwrap();
        let sqrt_b = get_sqrt_ratio_at_tick(100).unwrap();
        let amount1 = U256::from(1_000_000u64);
        let liq = liquidity_for_amount_1(sqrt_a, sqrt_b, amount1).unwrap();
        assert!(liq > 0);
    }

    #[test]
    fn test_max_liquidity_below_range() {
        let sqrt_current = get_sqrt_ratio_at_tick(-200).unwrap();
        let sqrt_lower = get_sqrt_ratio_at_tick(-100).unwrap();
        let sqrt_upper = get_sqrt_ratio_at_tick(100).unwrap();
        let liq = max_liquidity_for_amounts(
            sqrt_current,
            sqrt_lower,
            sqrt_upper,
            U256::from(1_000_000u64),
            U256::from(1_000_000u64),
        )
        .unwrap();
        // Below range: only amount0 matters.
        let liq0 =
            liquidity_for_amount_0(sqrt_lower, sqrt_upper, U256::from(1_000_000u64)).unwrap();
        assert_eq!(liq, liq0);
    }

    #[test]
    fn test_max_liquidity_above_range() {
        let sqrt_current = get_sqrt_ratio_at_tick(200).unwrap();
        let sqrt_lower = get_sqrt_ratio_at_tick(-100).unwrap();
        let sqrt_upper = get_sqrt_ratio_at_tick(100).unwrap();
        let liq = max_liquidity_for_amounts(
            sqrt_current,
            sqrt_lower,
            sqrt_upper,
            U256::from(1_000_000u64),
            U256::from(1_000_000u64),
        )
        .unwrap();
        let liq1 =
            liquidity_for_amount_1(sqrt_lower, sqrt_upper, U256::from(1_000_000u64)).unwrap();
        assert_eq!(liq, liq1);
    }

    #[test]
    fn test_same_sqrt_price() {
        let sqrt = get_sqrt_ratio_at_tick(0).unwrap();
        let liq = liquidity_for_amount_0(sqrt, sqrt, U256::from(1_000_000u64)).unwrap();
        assert_eq!(liq, 0);
    }

    #[test]
    fn fuzz_liquidity_for_amount_0_no_panic() {
        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(crate::tick_math::MIN_TICK..crate::tick_math::MAX_TICK);
            let tick_b: i32 = rng.random_range((tick_a + 1)..=crate::tick_math::MAX_TICK);
            let sqrt_a = get_sqrt_ratio_at_tick(tick_a).unwrap();
            let sqrt_b = get_sqrt_ratio_at_tick(tick_b).unwrap();
            let amount: u64 = rng.random_range(1..=u64::MAX);
            // Must not panic. Result can be Ok (including 0 from truncation)
            // or Err (overflow), both are valid.
            let _ = liquidity_for_amount_0(sqrt_a, sqrt_b, U256::from(amount));
        }
    }

    #[test]
    fn fuzz_liquidity_for_amount_1_no_panic() {
        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(crate::tick_math::MIN_TICK..crate::tick_math::MAX_TICK);
            let tick_b: i32 = rng.random_range((tick_a + 1)..=crate::tick_math::MAX_TICK);
            let sqrt_a = get_sqrt_ratio_at_tick(tick_a).unwrap();
            let sqrt_b = get_sqrt_ratio_at_tick(tick_b).unwrap();
            let amount: u64 = rng.random_range(1..=u64::MAX);
            // Must not panic. Ok (including 0) or Err (overflow) are both valid.
            let _ = liquidity_for_amount_1(sqrt_a, sqrt_b, U256::from(amount));
        }
    }

    #[test]
    fn fuzz_max_liquidity_for_amounts_no_panic() {
        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(crate::tick_math::MIN_TICK..crate::tick_math::MAX_TICK);
            let tick_b: i32 = rng.random_range((tick_a + 1)..=crate::tick_math::MAX_TICK);
            // Current tick can be anywhere in the valid range.
            let tick_current: i32 =
                rng.random_range(crate::tick_math::MIN_TICK..crate::tick_math::MAX_TICK);

            let sqrt_a = get_sqrt_ratio_at_tick(tick_a).unwrap();
            let sqrt_b = get_sqrt_ratio_at_tick(tick_b).unwrap();
            let sqrt_current = get_sqrt_ratio_at_tick(tick_current).unwrap();
            let amount0: u64 = rng.random_range(1..=u64::MAX);
            let amount1: u64 = rng.random_range(1..=u64::MAX);

            // Must not panic regardless of position relative to range.
            let _ = max_liquidity_for_amounts(
                sqrt_current,
                sqrt_a,
                sqrt_b,
                U256::from(amount0),
                U256::from(amount1),
            );
        }
    }

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

        // liquidity_for_amount_0(a, b, amt) == liquidity_for_amount_0(b, a, amt)
        let mut rng = test_rng();
        for _ in 0..1000 {
            let tick_a: i32 =
                rng.random_range(crate::tick_math::MIN_TICK..crate::tick_math::MAX_TICK);
            let tick_b: i32 = rng.random_range((tick_a + 1)..=crate::tick_math::MAX_TICK);
            let sqrt_a = get_sqrt_ratio_at_tick(tick_a).unwrap();
            let sqrt_b = get_sqrt_ratio_at_tick(tick_b).unwrap();
            let amount: u64 = rng.random_range(1..=1_000_000_000u64);

            let liq_ab = liquidity_for_amount_0(sqrt_a, sqrt_b, U256::from(amount));
            let liq_ba = liquidity_for_amount_0(sqrt_b, sqrt_a, U256::from(amount));
            match (liq_ab, liq_ba) {
                (Ok(a), Ok(b)) => assert_eq!(
                    a, b,
                    "ordering should not matter: tick_a={tick_a}, tick_b={tick_b}, amount={amount}"
                ),
                (Err(_), Err(_)) => {} // Both overflow — symmetric.
                _ => panic!(
                    "one succeeded, one failed: tick_a={tick_a}, tick_b={tick_b}, amount={amount}"
                ),
            }
        }
    }
}