sequence-algo-sdk 0.4.0

//! Pure AMM/CLMM pricing math for algo strategies.
//!
//! These functions let strategies compute swap output directly from raw pool
//! state (reserves, sqrt_price, tick, liquidity) rather than relying on
//! pre-computed synthetic order books.
//!
//! All functions are `no_std` compatible, use only integer arithmetic (no f64),
//! and have no system dependencies. Safe for WASM execution.
//!
//! ## Usage
//!
//! ```rust,ignore
//! use algo_sdk::amm_math;
//!
//! // V2 constant product (Uniswap V2, Raydium V4/CPMM)
//! let output = amm_math::v2_swap_output(amount_in, reserve_in, reserve_out, 3000);
//!
//! // V3/CLMM single-tick quote (Uniswap V3, Raydium CLMM, Orca Whirlpool)
//! let output = amm_math::v3_quote_single_tick(amount_in, sqrt_price_x64, liquidity, 3000);
//! ```

use crate::PoolAmm;

// ═════════════════════════════════════════════════════════════════════════════
// V2 Constant Product
// ═════════════════════════════════════════════════════════════════════════════

/// Compute output amount for a V2 constant-product swap.
///
/// Formula: `output = (input * (1_000_000 - fee_ppm) * reserve_out)
///                     / (reserve_in * 1_000_000 + input * (1_000_000 - fee_ppm))`
///
/// - `amount_in`: input token amount (raw units, NOT 1e8/1e9 scaled)
/// - `reserve_in`: pool reserve of the input token
/// - `reserve_out`: pool reserve of the output token
/// - `fee_ppm`: swap fee in parts-per-million (3000 = 0.3%)
///
/// Returns `None` if any input is zero or arithmetic overflows.
pub fn v2_swap_output(
    amount_in: u128,
    reserve_in: u128,
    reserve_out: u128,
    fee_ppm: u32,
) -> Option<u128> {
    if amount_in == 0 || reserve_in == 0 || reserve_out == 0 {
        return None;
    }
    if fee_ppm >= 1_000_000 {
        return None;
    }

    let fee_factor = (1_000_000 - fee_ppm) as u128;
    let amount_in_with_fee = amount_in.checked_mul(fee_factor)?;
    let denominator = reserve_in
        .checked_mul(1_000_000)?
        .checked_add(amount_in_with_fee)?;

    if denominator == 0 {
        return None;
    }
    // Use mul_div for 256-bit intermediate: (amount_in_with_fee * reserve_out) / denominator
    mul_div(amount_in_with_fee, reserve_out, denominator)
}

/// Compute the required input amount for a desired V2 output (inverse swap).
///
/// Returns `None` if the desired output exceeds reserves or inputs are zero.
pub fn v2_swap_input(
    amount_out: u128,
    reserve_in: u128,
    reserve_out: u128,
    fee_ppm: u32,
) -> Option<u128> {
    if amount_out == 0 || reserve_in == 0 || amount_out >= reserve_out {
        return None;
    }
    if fee_ppm >= 1_000_000 {
        return None;
    }

    let fee_factor = (1_000_000 - fee_ppm) as u128;
    let denominator = (reserve_out - amount_out).checked_mul(fee_factor)?;

    if denominator == 0 {
        return None;
    }
    // Use mul_div for 256-bit intermediate: ceil(reserve_in * amount_out * 1_000_000 / denominator)
    // amount_out * 1_000_000 is safe for any amount_out < 2^108 (all practical values)
    let amount_out_scaled = amount_out.checked_mul(1_000_000)?;
    let base = mul_div(reserve_in, amount_out_scaled, denominator)?;
    // +1 for ceiling: ensures the returned input is always sufficient
    Some(base + 1)
}

// ═════════════════════════════════════════════════════════════════════════════
// V3/CLMM Concentrated Liquidity
// ═════════════════════════════════════════════════════════════════════════════

/// Q64.64 fixed-point unit (2^64).
const Q64: u128 = 1u128 << 64;

/// Compute output for a V3/CLMM swap within a single tick range.
///
/// Accurate for trades small enough to stay within the current liquidity range.
/// For large trades that cross tick boundaries, this gives a lower bound.
///
/// - `amount_in`: input token amount (raw units)
/// - `sqrt_price_x64`: current sqrt(price) in Q64.64 format (from `PoolAmm.sqrt_price_x64`)
/// - `liquidity`: active liquidity at the current tick (from `PoolAmm.liquidity`)
/// - `fee_ppm`: swap fee in parts-per-million (3000 = 0.3%)
///
/// Returns `None` if inputs are zero or arithmetic overflows.
pub fn v3_quote_single_tick(
    amount_in: u128,
    sqrt_price_x64: u128,
    liquidity: u128,
    fee_ppm: u32,
) -> Option<u128> {
    if liquidity == 0 || sqrt_price_x64 == 0 || amount_in == 0 {
        return None;
    }
    if fee_ppm >= 1_000_000 {
        return None;
    }

    // Apply fee
    let amount_after_fee = amount_in * (1_000_000 - fee_ppm as u128) / 1_000_000;

    // Price = (sqrt_price_x64 / 2^64)^2 = sqrt_price_x64^2 / 2^128
    // For token0 → token1: output ≈ input * price
    // For token1 → token0: output ≈ input / price
    //
    // We compute: output = amount * (sqrt_price^2 / Q64^2)
    // Split to avoid overflow: (amount * sqrt_price / Q64) * (sqrt_price / Q64)
    let price_ratio = sqrt_price_x64;

    // output = amount_after_fee * price_ratio^2 / Q64^2
    // = (amount_after_fee * price_ratio / Q64) * price_ratio / Q64
    let step1 = mul_div(amount_after_fee, price_ratio, Q64)?;
    let result = mul_div(step1, price_ratio, Q64)?;

    if result > 0 { Some(result) } else { None }
}

/// Convert a tick index to sqrt_price in Q64.64 fixed-point.
///
/// `sqrt(1.0001^tick) = 1.00005^tick` — uses bit-decomposition with precomputed
/// magic constants in Q64.64 format.
///
/// Each constant is `round(1.00005^(2^bit) * 2^64)`.
///
/// Valid range: `-443636..=443636` (covers all practical price ranges).
pub fn tick_to_sqrt_price_x64(tick: i32) -> u128 {
    let abs_tick = tick.unsigned_abs();

    // Start at Q64 (= 1.0 in Q64.64)
    let mut ratio: u128 = Q64;

    // Precomputed: MAGIC[i] = round(sqrt(1.0001)^(2^i) * 2^64)
    // sqrt(1.0001) ≈ 1.000049998750062...
    // Only bits 0-18 needed: 2^19 = 524288 > max practical tick (443636).
    const MAGIC: [u128; 19] = [
        18447666387855958016,   // bit 0:  sqrt(1.0001)^1
        18448588748116918272,   // bit 1:  sqrt(1.0001)^2
        18450433606991728640,   // bit 2:  sqrt(1.0001)^4
        18454123878217453568,   // bit 3:  sqrt(1.0001)^8
        18461506635089977344,   // bit 4:  sqrt(1.0001)^16
        18476281010653851648,   // bit 5:  sqrt(1.0001)^32
        18505865242158133248,   // bit 6:  sqrt(1.0001)^64
        18565175891880198144,   // bit 7:  sqrt(1.0001)^128
        18684368066214465536,   // bit 8:  sqrt(1.0001)^256
        18925053041274802176,   // bit 9:  sqrt(1.0001)^512
        19415764168675909632,   // bit 10: sqrt(1.0001)^1024
        20435687552629014528,   // bit 11: sqrt(1.0001)^2048
        22639080592215080960,   // bit 12: sqrt(1.0001)^4096
        27784196929975758848,   // bit 13: sqrt(1.0001)^8192
        41848122137926787072,   // bit 14: sqrt(1.0001)^16384
        94936283577910951936,   // bit 15: sqrt(1.0001)^32768
        488590176324437606400,  // bit 16: sqrt(1.0001)^65536
        12941056668150515367936, // bit 17: sqrt(1.0001)^131072
        9078618265592131460530176, // bit 18: sqrt(1.0001)^262144
    ];

    for (i, &magic) in MAGIC.iter().enumerate() {
        if abs_tick & (1u32 << i) != 0 {
            ratio = mul_shift(ratio, magic);
        }
    }

    if tick < 0 {
        // For negative ticks: invert. sqrt_price(−t) = 1/sqrt_price(t) in Q64.64
        // result = Q64^2 / ratio = 2^128 / ratio
        if ratio == 0 { return 0; }
        // 2^128 doesn't fit u128, use: 2^128 / r = (2^128 - 1) / r + correction
        u128::MAX / ratio
    } else {
        ratio
    }
}

/// Convert sqrt_price in Q64.64 to a tick index.
///
/// Inverse of `tick_to_sqrt_price_x64`. Returns the largest tick `t` such that
/// `tick_to_sqrt_price_x64(t) <= sqrt_price_x64`.
pub fn sqrt_price_x64_to_tick(sqrt_price_x64: u128) -> i32 {
    if sqrt_price_x64 == 0 {
        return i32::MIN;
    }

    // Binary search for the tick
    let mut lo: i32 = -443636;
    let mut hi: i32 = 443636;

    while lo < hi {
        let mid = lo + (hi - lo + 1) / 2;
        if tick_to_sqrt_price_x64(mid) <= sqrt_price_x64 {
            lo = mid;
        } else {
            hi = mid - 1;
        }
    }
    lo
}

// ═════════════════════════════════════════════════════════════════════════════
// Convenience: compute from PoolAmm directly
// ═════════════════════════════════════════════════════════════════════════════

/// Compute swap output for a given input amount and pool state.
///
/// Automatically dispatches to V2 or V3 math based on `pool.pool_type`.
/// Returns `(amount_out, effective_price_1e9)` or `None` if the pool state
/// is insufficient or the trade would fail.
///
/// - `pool`: raw AMM state from `PoolStateTable`
/// - `amount_in`: input token amount (raw units)
/// - `zero_for_one`: true if swapping token0→token1 (selling base)
pub fn quote_from_pool(
    pool: &PoolAmm,
    amount_in: u128,
    zero_for_one: bool,
) -> Option<u128> {
    if !pool.is_valid() {
        return None;
    }

    match pool.pool_type {
        super::pool_type::V2 => {
            let (reserve_in, reserve_out) = if zero_for_one {
                (pool.reserve0(), pool.reserve1())
            } else {
                (pool.reserve1(), pool.reserve0())
            };
            v2_swap_output(amount_in, reserve_in, reserve_out, pool.fee_ppm)
        }
        super::pool_type::V3_CLMM => {
            v3_quote_single_tick(amount_in, pool.sqrt_price_x64, pool.liquidity, pool.fee_ppm)
        }
        _ => None, // stableswap not yet supported
    }
}

// ═════════════════════════════════════════════════════════════════════════════
// Internal helpers
// ═════════════════════════════════════════════════════════════════════════════

/// Multiply and right-shift by 64 (for Q64.64 fixed-point chain multiplication).
/// Computes `(a * b) >> 64` using a 256-bit intermediate to avoid overflow.
#[inline(always)]
fn mul_shift(a: u128, b: u128) -> u128 {
    // (a * b) >> 64 = (a_hi*b_hi) << 64
    //               + a_hi*b_lo
    //               + a_lo*b_hi
    //               + (a_lo*b_lo) >> 64
    //
    // The lower 128 bits of the sum are our result.
    let a_lo = a & 0xFFFF_FFFF_FFFF_FFFF;
    let a_hi = a >> 64;
    let b_lo = b & 0xFFFF_FFFF_FFFF_FFFF;
    let b_hi = b >> 64;

    let hh = a_hi * b_hi;
    let hl = a_hi * b_lo;
    let lh = a_lo * b_hi;
    let ll_hi = (a_lo * b_lo) >> 64;

    // Accumulate: result = hh << 64 + hl + lh + ll_hi
    // hl and lh each fit in u128. Their sum might carry.
    let mid = hl.wrapping_add(lh).wrapping_add(ll_hi);

    // hh << 64 adds to the upper 64 bits of our result
    // mid contributes to both halves
    (hh << 64).wrapping_add(mid)
}

/// Full 128×128 → 256-bit multiplication.
/// Returns `(hi, lo)` where the product = `hi * 2^128 + lo`.
#[inline(always)]
fn full_mul_u128(a: u128, b: u128) -> (u128, u128) {
    let a_lo = a & 0xFFFF_FFFF_FFFF_FFFF;
    let a_hi = a >> 64;
    let b_lo = b & 0xFFFF_FFFF_FFFF_FFFF;
    let b_hi = b >> 64;

    let ll = a_lo * b_lo;
    let lh = a_lo * b_hi;
    let hl = a_hi * b_lo;
    let hh = a_hi * b_hi;

    // Combine: product = hh*2^128 + (lh+hl)*2^64 + ll
    let (mid, mid_carry) = lh.overflowing_add(hl);
    let (lo, lo_carry) = ll.overflowing_add(mid << 64);
    let hi = hh + (mid >> 64) + ((mid_carry as u128) << 64) + lo_carry as u128;

    (hi, lo)
}

/// Divide a 256-bit number `(hi:lo)` by a 128-bit divisor `d`.
/// Precondition: `hi < d` (guarantees the quotient fits in u128).
/// Uses binary long division — 128 iterations of shift-and-subtract.
#[inline(always)]
fn div_256_by_128(hi: u128, lo: u128, d: u128) -> u128 {
    debug_assert!(hi < d, "quotient would overflow u128");
    if hi == 0 {
        return lo / d;
    }

    let mut rem = hi;
    let mut quot: u128 = 0;

    for i in (0..128).rev() {
        let bit = (lo >> i) & 1;
        // rem = 2*rem + bit; may overflow u128 when rem >= 2^127
        let will_overflow = rem >> 127 != 0;
        rem = rem.wrapping_shl(1) | bit;

        if will_overflow || rem >= d {
            rem = rem.wrapping_sub(d);
            quot |= 1u128 << i;
        }
    }

    quot
}

/// Multiply then divide with exact u128 precision: `floor(a * b / c)`.
/// Uses 256-bit intermediate to avoid overflow.
#[inline(always)]
fn mul_div(a: u128, b: u128, c: u128) -> Option<u128> {
    if c == 0 {
        return None;
    }
    // Fast path: no overflow possible
    if let Some(product) = a.checked_mul(b) {
        return Some(product / c);
    }
    // Full 256-bit path
    let (hi, lo) = full_mul_u128(a, b);
    if hi >= c {
        return None; // quotient > u128::MAX
    }
    Some(div_256_by_128(hi, lo, c))
}

// ═════════════════════════════════════════════════════════════════════════════
// Tests
// ═════════════════════════════════════════════════════════════════════════════

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

    #[test]
    fn v2_basic_swap() {
        // 1 ETH into 100 ETH / 200K USDC pool, 0.3% fee (3000 ppm)
        let out = v2_swap_output(
            1_000_000_000_000_000_000,   // 1e18 (1 ETH)
            100_000_000_000_000_000_000, // 100 ETH
            200_000_000_000,             // 200K USDC (6 decimals)
            3000,                        // 0.3%
        );
        assert!(out.is_some());
        let o = out.unwrap();
        // Should be ~1994 USDC (slightly less than 2000 due to fee + price impact)
        assert!(o > 1_970_000_000 && o < 2_000_000_000, "got {}", o);
    }

    #[test]
    fn v2_zero_inputs() {
        assert!(v2_swap_output(0, 100, 200, 3000).is_none());
        assert!(v2_swap_output(100, 0, 200, 3000).is_none());
        assert!(v2_swap_output(100, 200, 0, 3000).is_none());
    }

    #[test]
    fn v2_inverse_round_trip() {
        let reserve_in = 1_000_000_000u128;
        let reserve_out = 2_000_000_000u128;
        let amount_in = 10_000_000u128;
        let fee_ppm = 3000u32;

        let out = v2_swap_output(amount_in, reserve_in, reserve_out, fee_ppm).unwrap();
        let needed = v2_swap_input(out, reserve_in, reserve_out, fee_ppm).unwrap();

        // Due to ceiling division, needed >= amount_in
        assert!(needed >= amount_in);
        // But should be close (within 1 unit due to rounding)
        assert!(needed <= amount_in + 1, "needed={} vs in={}", needed, amount_in);
    }

    #[test]
    fn tick_to_price_zero() {
        let p = tick_to_sqrt_price_x64(0);
        // tick 0 = price 1.0, so sqrt_price = 1.0 in Q64.64 = 2^64
        assert_eq!(p, Q64);
    }

    #[test]
    fn tick_positive_gives_higher_price() {
        let p0 = tick_to_sqrt_price_x64(0);
        let p1 = tick_to_sqrt_price_x64(100);
        assert!(p1 > p0, "p1={}, p0={}", p1, p0);
    }

    #[test]
    fn tick_negative_gives_lower_price() {
        let p0 = tick_to_sqrt_price_x64(0);
        let pn = tick_to_sqrt_price_x64(-100);
        assert!(pn < p0, "pn={}, p0={}", pn, p0);
    }

    #[test]
    fn tick_round_trip() {
        for tick in [-1000, -100, -1, 0, 1, 100, 1000] {
            let price = tick_to_sqrt_price_x64(tick);
            let recovered = sqrt_price_x64_to_tick(price);
            assert_eq!(
                recovered, tick,
                "tick={} -> price={} -> recovered={}",
                tick, price, recovered
            );
        }
    }
}