wp-evm-v3-core 0.1.1

//! Native swap loop for v3-style concentrated-liquidity pools.
//!
//! Direct port of the pure-compute portion of Uniswap V3's
//! `UniswapV3Pool.sol::swap()`, built on top of `wp-evm-amm-math`
//! primitives (`compute_swap_step`, `tick_math::*`).
//!
//! Replaces the previous `uniswap-v3-sdk`-backed quoting path (removed
//! in R3 Slice D). The function is `pub(crate)` — internal to v3-core
//! per YAGNI. Callers go through
//! `quote::{exact_in,exact_out}_with_fee_fn`.
//!
//! Both exact-in (positive `amount_specified`) and exact-out (negative
//! `amount_specified`) go through the same loop — the sign distinguishes
//! them, per `UniswapV3Pool.sol::swap()`.

use alloy_primitives::{I256, U256};
use thiserror::Error;
use wp_evm_amm_math::{
    swap_math::compute_swap_step,
    tick_math::{
        get_sqrt_ratio_at_tick, get_tick_at_sqrt_ratio, MAX_SQRT_RATIO, MAX_TICK, MIN_SQRT_RATIO,
        MIN_TICK,
    },
    AmmMathError,
};

use crate::data::{PoolState, TickInfo};

/// Result of a completed swap. All values are unsigned; sign convention is
/// implicit in the swap direction (caller knows `zero_for_one`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct SwapResult {
    /// Actual input consumed (includes protocol fee). May be less than
    /// `amount_specified` for exact-in swaps that hit `sqrt_price_limit_x96`.
    pub amount_in: U256,
    /// Output produced for the consumed input.
    pub amount_out: U256,
    /// Pool's sqrt price after the swap.
    pub sqrt_price_x96_after: U256,
}

#[derive(Debug, Error)]
pub(crate) enum SwapError {
    #[error("amount_specified must be non-zero")]
    ZeroAmount,
    #[error("sqrt_price_limit_x96 invalid (must be on the correct side of current price and within [MIN_SQRT_RATIO, MAX_SQRT_RATIO])")]
    InvalidPriceLimit,
    #[error("ran out of liquidity before satisfying amount_specified")]
    InsufficientLiquidity,
    #[error("amm-math: {0}")]
    Math(#[from] AmmMathError),
    /// Invariant violation inside the swap loop. Used for conditions the
    /// surrounding code proves cannot occur (e.g. `i128::MIN.checked_neg()`
    /// on a `liquidity_net` field, or unreachable guard paths). Kept loud
    /// rather than `unreachable!()` so production builds surface the bug as
    /// an error rather than a panic.
    #[error("swap loop internal invariant violated: {0}")]
    Internal(&'static str),
}

/// Minimum valid `sqrt_price_limit_x96` for `zero_for_one` swaps (price
/// decreases). Equal to `MIN_SQRT_RATIO + 1` — one above the protocol
/// minimum, matching Solidity's `TickMath.MIN_SQRT_RATIO + 1` convention.
pub(crate) fn min_sqrt_ratio_plus_one() -> U256 {
    MIN_SQRT_RATIO + U256::from(1u64)
}

/// Maximum valid `sqrt_price_limit_x96` for `one_for_zero` swaps (price
/// increases). Equal to `MAX_SQRT_RATIO - 1`.
pub(crate) fn max_sqrt_ratio_minus_one() -> U256 {
    MAX_SQRT_RATIO - U256::from(1u64)
}

/// Apply an `i128` delta to a `u128` liquidity value with overflow checking.
///
/// Mirrors Solidity's `LiquidityMath.addDelta`. Defined here rather than in
/// `wp-evm-amm-math` per YAGNI — the swap loop is its only current caller.
fn add_delta(x: u128, y: i128) -> core::result::Result<u128, SwapError> {
    if y < 0 {
        let neg = y.unsigned_abs();
        x.checked_sub(neg).ok_or(SwapError::InsufficientLiquidity)
    } else {
        x.checked_add(y as u128).ok_or(SwapError::Math(AmmMathError::LiquidityOverflow))
    }
}

/// Locate the next initialized tick in the direction of the swap.
///
/// `state.ticks` is assumed sorted ascending by `tick` (see
/// `data::PoolState::ticks` contract). O(log n) via `partition_point` on
/// the tick coordinate — we convert `current_sqrt_x96` to its floor tick
/// once and then binary-search `t.tick`, rather than recomputing
/// `get_sqrt_ratio_at_tick` per candidate.
///
/// Returns `None` if no initialized tick exists in the requested direction.
fn next_initialized_tick(
    ticks: &[TickInfo],
    current_sqrt_x96: U256,
    zero_for_one: bool,
) -> core::result::Result<Option<&TickInfo>, SwapError> {
    if ticks.is_empty() {
        return Ok(None);
    }

    // `get_tick_at_sqrt_ratio` returns the greatest tick `t` such that
    // `get_sqrt_ratio_at_tick(t) <= current_sqrt_x96`. Combined with whether
    // current is exactly on a tick boundary, we can express both search
    // predicates purely in terms of `t.tick`.
    let cur_tick = get_tick_at_sqrt_ratio(current_sqrt_x96)?;
    let cur_at_tick = get_sqrt_ratio_at_tick(cur_tick)?;
    let on_tick_boundary = cur_at_tick == current_sqrt_x96;

    if zero_for_one {
        // Largest tick with sqrt(tick) < current.
        // - on boundary: strict inequality requires t.tick < cur_tick.
        // - off boundary: sqrt(cur_tick) < current, so t.tick <= cur_tick qualifies.
        let idx = if on_tick_boundary {
            ticks.partition_point(|t| t.tick < cur_tick)
        } else {
            ticks.partition_point(|t| t.tick <= cur_tick)
        };
        Ok(if idx == 0 { None } else { Some(&ticks[idx - 1]) })
    } else {
        // Smallest tick with sqrt(tick) > current. Since
        // sqrt(cur_tick) <= current < sqrt(cur_tick + 1), the condition
        // reduces to t.tick > cur_tick regardless of boundary coincidence.
        let idx = ticks.partition_point(|t| t.tick <= cur_tick);
        Ok(ticks.get(idx))
    }
}

/// Convenience result alias for this module.
type Result<T> = core::result::Result<T, SwapError>;

/// Native swap loop — port of `UniswapV3Pool.sol::swap()`'s pure-compute path.
///
/// - `zero_for_one = true`: token0 in → token1 out, price moves DOWN.
/// - `amount_specified > 0`: exact-in (caller specifies input).
/// - `amount_specified < 0`: exact-out (caller specifies output magnitude).
///   Both directions are live — exercised by
///   `quote::{exact_in,exact_out}_with_fee_fn` and by the native-vs-SDK
///   parity harness in `quote::parity`.
///
/// `sqrt_price_limit_x96` must be:
/// - strictly less than current sqrt price for `zero_for_one = true`;
/// - strictly greater than current sqrt price for `zero_for_one = false`;
/// - within `(MIN_SQRT_RATIO, MAX_SQRT_RATIO)` exclusive.
///
/// `fee_pips` is the fee in hundredths of a basis point (e.g. 3000 for 0.30%).
pub(crate) fn swap(
    state: &PoolState,
    zero_for_one: bool,
    amount_specified: I256,
    sqrt_price_limit_x96: U256,
    fee_pips: u32,
) -> Result<SwapResult> {
    if amount_specified.is_zero() {
        return Err(SwapError::ZeroAmount);
    }

    // Guard the price limit. Mirrors UniswapV3Pool::swap require() block.
    if zero_for_one {
        if sqrt_price_limit_x96 >= state.sqrt_price_x96 || sqrt_price_limit_x96 <= MIN_SQRT_RATIO {
            return Err(SwapError::InvalidPriceLimit);
        }
    } else if sqrt_price_limit_x96 <= state.sqrt_price_x96 || sqrt_price_limit_x96 >= MAX_SQRT_RATIO
    {
        return Err(SwapError::InvalidPriceLimit);
    }

    let exact_in = !amount_specified.is_negative();

    // UniswapV3Pool.sol lines ~620–720. Loop state mirrors the Solidity
    // locals with 1:1 naming.
    let mut amount_specified_remaining: I256 = amount_specified;
    let mut amount_calculated: I256 = I256::ZERO;
    let mut sqrt_price_x96: U256 = state.sqrt_price_x96;
    let mut liquidity: u128 = state.liquidity;

    // Fallback boundary for the "no more initialized ticks in direction"
    // case: clamp to the sqrt price at the protocol tick boundary.
    let boundary_sqrt = if zero_for_one {
        get_sqrt_ratio_at_tick(MIN_TICK)?
    } else {
        get_sqrt_ratio_at_tick(MAX_TICK)?
    };

    // Loop invariant: `sqrt_price_x96` stays strictly inside
    // `(MIN_SQRT_RATIO, MAX_SQRT_RATIO)`. Maintained by the price-limit
    // guard (caller limit is strictly inside that range) and the
    // target-clamp step: `target_sqrt` is always at least as far from the
    // protocol boundary as the caller limit. This invariant is required
    // by `next_initialized_tick`, which calls `get_tick_at_sqrt_ratio`
    // (errors outside that range).
    while !amount_specified_remaining.is_zero() && sqrt_price_x96 != sqrt_price_limit_x96 {
        // 1. Find the next initialized tick in-direction (or fall back to
        //    the protocol boundary sqrt). This is the step's candidate target.
        let next_tick_opt = next_initialized_tick(&state.ticks, sqrt_price_x96, zero_for_one)?;
        let next_tick_sqrt = match next_tick_opt {
            Some(t) => get_sqrt_ratio_at_tick(t.tick)?,
            None => boundary_sqrt,
        };

        // 2. Clamp the target to whichever is "closer to current":
        //    caller's limit, or the next tick sqrt. For zero_for_one (price
        //    going down), "closer to current" means LARGER.
        let target_sqrt = if zero_for_one {
            core::cmp::max(next_tick_sqrt, sqrt_price_limit_x96)
        } else {
            core::cmp::min(next_tick_sqrt, sqrt_price_limit_x96)
        };

        // 3. Execute one swap step against the current (sqrt, liquidity) range.
        //
        //    Invariant: liquidity==0 combined with a non-zero amount would
        //    loop forever (step returns zeros without advancing the price).
        //    Short-circuit to InsufficientLiquidity instead.
        if liquidity == 0 {
            return Err(SwapError::InsufficientLiquidity);
        }
        let step = compute_swap_step(
            sqrt_price_x96,
            target_sqrt,
            liquidity,
            amount_specified_remaining,
            fee_pips,
        )?;

        // 4. Accumulate into the I256 counters. Solidity uses `unchecked`;
        //    we use `checked_*` + ok_or panic-equivalent since the
        //    per-step amounts are bounded by pool liquidity which is itself
        //    bounded by int256. Any overflow here is a real bug, not a
        //    pool-state concern — surface it rather than wrap silently.
        //    Ref: UniswapV3Pool.sol L664–L679.
        let step_in_with_fee_u: U256 = step
            .amount_in
            .checked_add(step.fee_amount)
            .ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
        let step_in_with_fee: I256 = I256::try_from(step_in_with_fee_u)
            .map_err(|_| SwapError::Math(AmmMathError::MulDivOverflow))?;
        let step_out: I256 = I256::try_from(step.amount_out)
            .map_err(|_| SwapError::Math(AmmMathError::MulDivOverflow))?;

        if exact_in {
            amount_specified_remaining = amount_specified_remaining
                .checked_sub(step_in_with_fee)
                .ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
            amount_calculated = amount_calculated
                .checked_sub(step_out)
                .ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
        } else {
            // Exact-out: remaining is negative; we add step_out (positive)
            // to move it toward zero. amount_calculated accumulates the
            // positive input + fee.
            amount_specified_remaining = amount_specified_remaining
                .checked_add(step_out)
                .ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
            amount_calculated = amount_calculated
                .checked_add(step_in_with_fee)
                .ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
        }

        // 5. Advance the pool price.
        sqrt_price_x96 = step.sqrt_ratio_next_x96;

        // 6. If we landed exactly on the next initialized tick (not just a
        //    clamp to the caller's price limit or to the protocol boundary),
        //    cross it: apply the liquidity_net delta with a sign flip for
        //    zero_for_one. Solidity: `liquidityNet = zeroForOne ? -liquidityNet : liquidityNet`.
        if sqrt_price_x96 == next_tick_sqrt {
            if let Some(next_tick) = next_tick_opt {
                let delta = if zero_for_one {
                    // Negation overflow is only reachable at i128::MIN;
                    // real pools never encode |liquidity_net| = 2^127.
                    // Surface as Internal rather than panic so the signal
                    // survives in production if the invariant ever breaks.
                    next_tick
                        .liquidity_net
                        .checked_neg()
                        .ok_or(SwapError::Internal("liquidity_net negation overflow (i128::MIN)"))?
                } else {
                    next_tick.liquidity_net
                };
                liquidity = add_delta(liquidity, delta)?;
            }
            // If next_tick_opt was None we hit the protocol boundary —
            // no tick to cross. Given the input guards on
            // sqrt_price_limit_x96, target_sqrt is clamped to the caller
            // limit (never boundary_sqrt), so the loop exits naturally
            // via the while condition and the remainder check below
            // promotes any leftover amount to InsufficientLiquidity.
        }
    }

    // If we bailed at the protocol boundary with remainder, that's genuine
    // "pool cannot provide this much". If we bailed at the caller's explicit
    // sqrt_price_limit_x96, that's a valid partial fill — return Ok.
    if !amount_specified_remaining.is_zero() && sqrt_price_x96 != sqrt_price_limit_x96 {
        return Err(SwapError::InsufficientLiquidity);
    }

    // Translate signed accumulators back to unsigned (amount_in, amount_out).
    // Exact-in:  consumed = amount_specified - remaining   (positive I256)
    //            produced = -amount_calculated             (positive I256)
    // Exact-out: produced = amount_specified - remaining   (negative I256, magnitude is output)
    //            consumed = amount_calculated              (positive I256)
    let (amount_in_i, amount_out_i) = if exact_in {
        let consumed = amount_specified
            .checked_sub(amount_specified_remaining)
            .ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
        let produced =
            amount_calculated.checked_neg().ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
        (consumed, produced)
    } else {
        let produced_neg = amount_specified
            .checked_sub(amount_specified_remaining)
            .ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
        let produced =
            produced_neg.checked_neg().ok_or(SwapError::Math(AmmMathError::MulDivOverflow))?;
        (amount_calculated, produced)
    };

    let amount_in: U256 = if amount_in_i.is_negative() {
        return Err(SwapError::Math(AmmMathError::MulDivOverflow));
    } else {
        amount_in_i.into_raw()
    };
    let amount_out: U256 = if amount_out_i.is_negative() {
        return Err(SwapError::Math(AmmMathError::MulDivOverflow));
    } else {
        amount_out_i.into_raw()
    };

    Ok(SwapResult { amount_in, amount_out, sqrt_price_x96_after: sqrt_price_x96 })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::data::{PoolState, TickInfo};
    use alloy_primitives::address;

    /// Mirror of quote::tests::fixture_usdc_weth_03 — duplicated here so
    /// swap-loop unit tests stay decoupled from quote.rs test layout.
    fn fixture_usdc_weth_03() -> PoolState {
        let sqrt_price_x96: U256 =
            U256::from_str_radix("3543191142285914205922034323214", 10).unwrap();
        PoolState {
            token0: address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"),
            token1: address!("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"),
            fee: 3000,
            tick_spacing: 60,
            sqrt_price_x96,
            liquidity: 2_000_000_000_000_000_000_000u128,
            tick: 76012,
            ticks: vec![
                TickInfo {
                    tick: 74940,
                    liquidity_net: 1_000_000_000_000_000_000_000i128,
                    liquidity_gross: 1_000_000_000_000_000_000_000u128,
                },
                TickInfo {
                    tick: 75960,
                    liquidity_net: 1_000_000_000_000_000_000_000i128,
                    liquidity_gross: 1_000_000_000_000_000_000_000u128,
                },
                TickInfo {
                    tick: 76020,
                    liquidity_net: -2_000_000_000_000_000_000_000i128,
                    liquidity_gross: 2_000_000_000_000_000_000_000u128,
                },
            ],
        }
    }

    #[test]
    fn swap_no_tick_cross_partial_fill() {
        let s = fixture_usdc_weth_03();
        let amt = I256::try_from(U256::from(1_000u64)).unwrap();
        let r = swap(&s, true, amt, min_sqrt_ratio_plus_one(), 3000).unwrap();
        assert!(r.amount_in > U256::ZERO);
        assert!(r.amount_out > U256::ZERO);
        // Tiny input: price should move very little and stay well above tick 75960.
        let sqrt_at_75960 = U256::from_str_radix("3533845506420911390540068078527", 10).unwrap();
        assert!(r.sqrt_price_x96_after > sqrt_at_75960);
        assert!(r.sqrt_price_x96_after < s.sqrt_price_x96);
    }

    #[test]
    fn swap_crosses_one_tick() {
        // 1e18 USDC input crosses tick 75960 downward. After crossing,
        // active liquidity drops from 2e21 to 1e21 (per fixture).
        let s = fixture_usdc_weth_03();
        let amt = I256::try_from(U256::from(1_000_000_000_000_000_000u64)).unwrap();
        let r = swap(&s, true, amt, min_sqrt_ratio_plus_one(), 3000).unwrap();
        let sqrt_at_75960 = U256::from_str_radix("3533845506420911390540068078527", 10).unwrap();
        assert!(
            r.sqrt_price_x96_after < sqrt_at_75960,
            "should cross tick 75960, ended at {}",
            r.sqrt_price_x96_after
        );
        assert_eq!(r.amount_in, U256::from(1_000_000_000_000_000_000u64));
    }

    #[test]
    fn swap_clamps_at_price_limit() {
        // Set the limit to a sqrt price between current and tick 75960;
        // assert loop exits at limit and leaves remainder (amount_in <
        // amount_specified).
        let s = fixture_usdc_weth_03();
        let amt = I256::try_from(U256::from(1_000_000_000_000_000_000u64)).unwrap();
        // Pick a limit just above tick 75960 sqrt — ensures we clamp here,
        // don't cross.
        let custom_limit = U256::from_str_radix("3540000000000000000000000000000", 10).unwrap();
        assert!(custom_limit < s.sqrt_price_x96);
        let r = swap(&s, true, amt, custom_limit, 3000).unwrap();
        assert_eq!(r.sqrt_price_x96_after, custom_limit, "should clamp at caller-provided limit");
        assert!(
            r.amount_in < U256::from(1_000_000_000_000_000_000u64),
            "should leave input unconsumed: got {}",
            r.amount_in
        );
    }

    #[test]
    fn swap_empty_ticks_zero_for_one_errors() {
        // liquidity=0 + empty ticks: any nonzero amount must fail with
        // InsufficientLiquidity (no liquidity to swap against).
        let mut s = fixture_usdc_weth_03();
        s.ticks.clear();
        s.liquidity = 0;
        let amt = I256::try_from(U256::from(1_000u64)).unwrap();
        let err = swap(&s, true, amt, min_sqrt_ratio_plus_one(), 3000).unwrap_err();
        assert!(matches!(err, SwapError::InsufficientLiquidity));
    }

    #[test]
    fn swap_direction_zero_for_one_decreases_price() {
        let s = fixture_usdc_weth_03();
        let amt = I256::try_from(U256::from(1_000_000u64)).unwrap();
        let r = swap(&s, true, amt, min_sqrt_ratio_plus_one(), 3000).unwrap();
        assert!(r.sqrt_price_x96_after < s.sqrt_price_x96);
    }

    #[test]
    fn swap_direction_one_for_zero_increases_price() {
        let s = fixture_usdc_weth_03();
        let amt = I256::try_from(U256::from(1_000_000_000_000_000u64)).unwrap(); // WETH
        let r = swap(&s, false, amt, max_sqrt_ratio_minus_one(), 3000).unwrap();
        assert!(r.sqrt_price_x96_after > s.sqrt_price_x96);
    }

    #[test]
    fn swap_exact_in_conserves_input_and_fee() {
        // For exact-in without hitting the price limit, amount_in must equal
        // amount_specified (since the loop consumes it all). For a partial
        // fill (price-limit hit), amount_in must be < amount_specified.
        let s = fixture_usdc_weth_03();
        let amt_u = U256::from(1_000_000u64);
        let amt = I256::try_from(amt_u).unwrap();
        let r = swap(&s, true, amt, min_sqrt_ratio_plus_one(), 3000).unwrap();
        assert_eq!(r.amount_in, amt_u, "exact-in should fully consume input");
    }

    #[test]
    fn swap_rejects_zero_amount() {
        let s = fixture_usdc_weth_03();
        let err = swap(&s, true, I256::ZERO, min_sqrt_ratio_plus_one(), 3000).unwrap_err();
        assert!(matches!(err, SwapError::ZeroAmount));
    }

    #[test]
    fn swap_exact_out_no_tick_cross() {
        // Small WETH output request: doesn't reach tick 75960. amount_out
        // must equal the requested magnitude; amount_in is derived.
        let s = fixture_usdc_weth_03();
        let req_out = U256::from(1_000_000_000u64); // 1e9 wei WETH
        let amt = I256::try_from(req_out).unwrap().checked_neg().unwrap();
        let r = swap(&s, true, amt, min_sqrt_ratio_plus_one(), 3000).unwrap();
        assert_eq!(r.amount_out, req_out);
        assert!(r.amount_in > U256::ZERO);
        assert!(r.sqrt_price_x96_after < s.sqrt_price_x96);
        let sqrt_at_75960 = U256::from_str_radix("3533845506420911390540068078527", 10).unwrap();
        assert!(r.sqrt_price_x96_after > sqrt_at_75960, "should not cross tick 75960");
    }

    #[test]
    fn swap_exact_out_crosses_one_tick() {
        // The current range [75960..76020] at L=2e21 holds roughly 2.36e20
        // WETH (L * delta_sqrt / Q96). Requesting 3e20 forces the loop to
        // cross tick 75960 downward; after crossing, liquidity drops to
        // 1e21 (wide range only) and the remaining ~6.4e19 is drawn from
        // [74940..75960] without hitting the next boundary.
        let s = fixture_usdc_weth_03();
        let req_out = U256::from_str_radix("300000000000000000000", 10).unwrap(); // 3e20
        let amt = I256::try_from(req_out).unwrap().checked_neg().unwrap();
        let r = swap(&s, true, amt, min_sqrt_ratio_plus_one(), 3000).unwrap();
        assert_eq!(r.amount_out, req_out, "exact-out must supply full request");
        let sqrt_at_75960 = U256::from_str_radix("3533845506420911390540068078527", 10).unwrap();
        assert!(
            r.sqrt_price_x96_after < sqrt_at_75960,
            "should cross tick 75960, ended at {}",
            r.sqrt_price_x96_after
        );
    }

    #[test]
    fn swap_rejects_invalid_price_limit() {
        let s = fixture_usdc_weth_03();
        let amt = I256::try_from(U256::from(1_000u64)).unwrap();
        // zero_for_one with limit ABOVE current is invalid.
        let bad_limit = s.sqrt_price_x96 + U256::from(1u64);
        let err = swap(&s, true, amt, bad_limit, 3000).unwrap_err();
        assert!(matches!(err, SwapError::InvalidPriceLimit));
    }
}