wp-solana-amm-math 0.1.1

//! Core swap step computation (Uniswap V3 / Orca / Raydium CLMM).
//!
//! All functions are pure math -- no RPC, no account types.

use ethnum::U256;

use crate::{
    fee_math::fee_amount_from_input,
    liquidity_math::{get_amount_0_delta, get_amount_1_delta},
    AmmMathError,
};

/// Result of a single swap step within one tick range.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SwapStepResult {
    /// The sqrt price after the step, in Q64.64.
    pub sqrt_price_next: u128,
    /// Amount of input token consumed (excluding fees).
    pub amount_in: u64,
    /// Amount of output token produced.
    pub amount_out: u64,
    /// Fee amount deducted from the input.
    pub fee_amount: u64,
}

/// Compute the next sqrt price when adding `amount_in` tokens as input.
///
/// # Q64.64 arithmetic
///
/// - **a_to_b** (adding token A, price decreases): `sqrt_price_next = L *
///   sqrt_price / (L + amount * sqrt_price)` all in Q64.64 via U256
///   intermediates.
///
/// - **b_to_a** (adding token B, price increases): `sqrt_price_next =
///   sqrt_price + amount * 2^64 / L`
pub fn get_next_sqrt_price_from_input(
    sqrt_price: u128,
    liquidity: u128,
    amount_in: u64,
    a_to_b: bool,
) -> Result<u128, AmmMathError> {
    if sqrt_price == 0 {
        return Err(AmmMathError::SqrtPriceOutOfRange(0));
    }
    if liquidity == 0 {
        return Err(AmmMathError::DivisionByZero);
    }

    if amount_in == 0 {
        return Ok(sqrt_price);
    }

    if a_to_b {
        // sqrt_price_next = (L << 64) * sqrt_price
        //                 / ((L << 64) + amount * sqrt_price)
        let l_shifted = U256::from(liquidity) << 64;
        let amount = U256::from(amount_in);
        let price = U256::from(sqrt_price);

        let numerator = l_shifted * price;
        let denominator = l_shifted + amount * price;

        if denominator == U256::ZERO {
            return Err(AmmMathError::DivisionByZero);
        }

        // Round UP so the resulting price does not undershoot the exact
        // value (matches Uniswap V3 / Orca Whirlpools reference).
        let result: U256 = (numerator + denominator - U256::from(1u8)) / denominator;
        if result > U256::from(u128::MAX) {
            return Err(AmmMathError::Overflow);
        }
        Ok(result.as_u128())
    } else {
        // sqrt_price_next = sqrt_price + (amount << 64) / L
        let delta = (U256::from(amount_in) << 64) / U256::from(liquidity);
        let result: U256 = U256::from(sqrt_price) + delta;
        if result > U256::from(u128::MAX) {
            return Err(AmmMathError::Overflow);
        }
        Ok(result.as_u128())
    }
}

/// Compute the next sqrt price when removing `amount_out` tokens as
/// output.
///
/// - **a_to_b** (removing token B, price decreases): `sqrt_price_next =
///   sqrt_price - ceil(amount * 2^64 / L)`
///
/// - **b_to_a** (removing token A, price increases): `sqrt_price_next = L *
///   sqrt_price / (L - amount * sqrt_price)` rounded up.
pub fn get_next_sqrt_price_from_output(
    sqrt_price: u128,
    liquidity: u128,
    amount_out: u64,
    a_to_b: bool,
) -> Result<u128, AmmMathError> {
    if sqrt_price == 0 {
        return Err(AmmMathError::SqrtPriceOutOfRange(0));
    }
    if liquidity == 0 {
        return Err(AmmMathError::DivisionByZero);
    }

    if amount_out == 0 {
        return Ok(sqrt_price);
    }

    if a_to_b {
        // Price decreases: removing token B (token 1).
        // delta = ceil((amount << 64) / L)
        let numerator = U256::from(amount_out) << 64;
        let denom = U256::from(liquidity);
        let delta = (numerator + denom - U256::ONE) / denom;

        let price: U256 = U256::from(sqrt_price);
        if delta >= price {
            return Err(AmmMathError::Overflow);
        }
        let next: U256 = price - delta;
        Ok(next.as_u128())
    } else {
        // Price increases: removing token A (token 0).
        // sqrt_price_next = ceil(L * sqrt_price / (L - amount * sqrt_price))
        let l = U256::from(liquidity);
        let price = U256::from(sqrt_price);
        let amount = U256::from(amount_out);

        let product = amount * price;
        let l_shifted = l << 64;
        if product >= l_shifted {
            return Err(AmmMathError::Overflow);
        }
        let denominator = l_shifted - product;
        let numerator = l_shifted * price;

        // Ceiling division
        let result: U256 = (numerator + denominator - U256::ONE) / denominator;
        if result > U256::from(u128::MAX) {
            return Err(AmmMathError::Overflow);
        }
        Ok(result.as_u128())
    }
}

/// Compute a single swap step within one liquidity range.
///
/// # Arguments
///
/// * `sqrt_price_current` - Current sqrt price in Q64.64.
/// * `sqrt_price_target` - Target sqrt price (next initialized tick).
/// * `liquidity` - Active liquidity in the range.
/// * `amount_remaining` - Remaining swap amount (input or output).
/// * `fee_rate_bps` - Fee rate in basis points (e.g. 30 = 0.3%).
/// * `by_amount_in` - `true` for ExactIn, `false` for ExactOut.
pub fn compute_swap_step(
    sqrt_price_current: u128,
    sqrt_price_target: u128,
    liquidity: u128,
    amount_remaining: u64,
    fee_rate_bps: u16,
    by_amount_in: bool,
) -> Result<SwapStepResult, AmmMathError> {
    let a_to_b = sqrt_price_current >= sqrt_price_target;

    if amount_remaining == 0 || liquidity == 0 {
        return Ok(SwapStepResult {
            sqrt_price_next: sqrt_price_current,
            amount_in: 0,
            amount_out: 0,
            fee_amount: 0,
        });
    }

    let sqrt_price_next;
    let amount_in;
    let amount_out;
    let fee_amount;

    if by_amount_in {
        // ExactIn: amount_remaining is the input amount (inclusive of fee).
        let fee_on_remaining = fee_amount_from_input(amount_remaining, fee_rate_bps)?;
        let amount_remaining_less_fee = amount_remaining.saturating_sub(fee_on_remaining);

        // Max amount consumable to reach the target price.
        let max_amount_in = if a_to_b {
            // Selling token A: compute delta_x to reach target.
            get_amount_0_delta(sqrt_price_target, sqrt_price_current, liquidity, true)?
        } else {
            // Selling token B: compute delta_y to reach target.
            get_amount_1_delta(sqrt_price_current, sqrt_price_target, liquidity, true)?
        };

        if amount_remaining_less_fee >= max_amount_in {
            // We can reach the target price.
            sqrt_price_next = sqrt_price_target;
            amount_in = max_amount_in;
        } else {
            // Partial fill: compute new price from available input.
            sqrt_price_next = get_next_sqrt_price_from_input(
                sqrt_price_current,
                liquidity,
                amount_remaining_less_fee,
                a_to_b,
            )?;
            // Recompute amount_in from the price delta (rounded up), but
            // cap at the available input to avoid rounding overshoot.
            let delta_in = if a_to_b {
                get_amount_0_delta(sqrt_price_next, sqrt_price_current, liquidity, true)?
            } else {
                get_amount_1_delta(sqrt_price_current, sqrt_price_next, liquidity, true)?
            };
            amount_in = delta_in.min(amount_remaining_less_fee);
        }

        // Compute output from the price movement.
        amount_out = if a_to_b {
            get_amount_1_delta(sqrt_price_next, sqrt_price_current, liquidity, false)?
        } else {
            get_amount_0_delta(sqrt_price_current, sqrt_price_next, liquidity, false)?
        };

        // Fee calculation.
        if sqrt_price_next != sqrt_price_target {
            // Didn't reach target: consume all remaining as fee + input.
            fee_amount = amount_remaining.saturating_sub(amount_in);
        } else {
            fee_amount = fee_amount_from_input(amount_in, fee_rate_bps)?;
        }
    } else {
        // ExactOut: amount_remaining is the desired output.
        let max_amount_out = if a_to_b {
            get_amount_1_delta(sqrt_price_target, sqrt_price_current, liquidity, false)?
        } else {
            get_amount_0_delta(sqrt_price_current, sqrt_price_target, liquidity, false)?
        };

        let capped_output = amount_remaining.min(max_amount_out);

        if capped_output == 0 {
            return Ok(SwapStepResult {
                sqrt_price_next: sqrt_price_current,
                amount_in: 0,
                amount_out: 0,
                fee_amount: 0,
            });
        }

        sqrt_price_next = if amount_remaining >= max_amount_out {
            sqrt_price_target
        } else {
            get_next_sqrt_price_from_output(sqrt_price_current, liquidity, capped_output, a_to_b)?
        };

        amount_out = capped_output;

        // Compute the input required for this price movement.
        amount_in = if a_to_b {
            get_amount_0_delta(sqrt_price_next, sqrt_price_current, liquidity, true)?
        } else {
            get_amount_1_delta(sqrt_price_current, sqrt_price_next, liquidity, true)?
        };

        fee_amount = fee_amount_from_input(amount_in, fee_rate_bps)?;
    }

    Ok(SwapStepResult { sqrt_price_next, amount_in, amount_out, fee_amount })
}

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

    const Q64: u128 = 1u128 << 64;

    // ------------------------------------------------------------------
    // get_next_sqrt_price_from_input
    // ------------------------------------------------------------------

    #[test]
    fn test_next_price_from_input_zero_amount() {
        let price = Q64;
        let result = get_next_sqrt_price_from_input(price, 1_000_000, 0, true).unwrap();
        assert_eq!(result, price);
    }

    #[test]
    fn test_next_price_from_input_a_to_b() {
        // Adding token A => price decreases
        let price = Q64; // price = 1.0
        let liquidity = 1_000_000u128;
        let result = get_next_sqrt_price_from_input(price, liquidity, 100, true).unwrap();
        assert!(result < price, "a_to_b should decrease price");
    }

    #[test]
    fn test_next_price_from_input_b_to_a() {
        // Adding token B => price increases
        let price = Q64;
        let liquidity = 1_000_000u128;
        let result = get_next_sqrt_price_from_input(price, liquidity, 100, false).unwrap();
        assert!(result > price, "b_to_a should increase price");
    }

    #[test]
    fn test_next_price_from_input_zero_liquidity() {
        assert!(get_next_sqrt_price_from_input(Q64, 0, 100, true).is_err());
    }

    #[test]
    fn test_next_price_from_input_zero_price() {
        assert!(get_next_sqrt_price_from_input(0, 1000, 100, true).is_err());
    }

    // ------------------------------------------------------------------
    // get_next_sqrt_price_from_output
    // ------------------------------------------------------------------

    #[test]
    fn test_next_price_from_output_zero_amount() {
        let result = get_next_sqrt_price_from_output(Q64, 1_000_000, 0, true).unwrap();
        assert_eq!(result, Q64);
    }

    #[test]
    fn test_next_price_from_output_a_to_b() {
        // Removing token B => price decreases
        let price = Q64;
        let liquidity = 1_000_000_000u128;
        let result = get_next_sqrt_price_from_output(price, liquidity, 100, true).unwrap();
        assert!(result < price, "a_to_b output should decrease price");
    }

    #[test]
    fn test_next_price_from_output_b_to_a() {
        // Removing token A => price increases
        let price = Q64;
        let liquidity = 1_000_000_000u128;
        let result = get_next_sqrt_price_from_output(price, liquidity, 100, false).unwrap();
        assert!(result > price, "b_to_a output should increase price");
    }

    #[test]
    fn test_next_price_from_output_excessive_amount_a_to_b() {
        // Trying to remove more token B than available should error
        let result = get_next_sqrt_price_from_output(Q64, 1, u64::MAX, true);
        assert!(result.is_err());
    }

    // ------------------------------------------------------------------
    // compute_swap_step
    // ------------------------------------------------------------------

    #[test]
    fn test_swap_step_zero_amount() {
        let result = compute_swap_step(Q64, Q64 / 2, 1_000_000, 0, 30, true).unwrap();
        assert_eq!(result.amount_in, 0);
        assert_eq!(result.amount_out, 0);
        assert_eq!(result.fee_amount, 0);
        assert_eq!(result.sqrt_price_next, Q64);
    }

    #[test]
    fn test_swap_step_zero_liquidity() {
        let result = compute_swap_step(Q64, Q64 / 2, 0, 1000, 30, true).unwrap();
        assert_eq!(result.amount_in, 0);
        assert_eq!(result.amount_out, 0);
    }

    #[test]
    fn test_swap_step_exact_in_a_to_b_reaches_target() {
        // Set up: large amount, should reach target price
        let price_current = tick_to_sqrt_price_x64(100).unwrap();
        let price_target = tick_to_sqrt_price_x64(0).unwrap();
        let liquidity = 10_000_000_000u128;

        let result =
            compute_swap_step(price_current, price_target, liquidity, u64::MAX, 30, true).unwrap();

        assert_eq!(result.sqrt_price_next, price_target, "should reach target price");
        assert!(result.amount_in > 0);
        assert!(result.amount_out > 0);
        assert!(result.fee_amount > 0);
    }

    #[test]
    fn test_swap_step_exact_in_a_to_b_partial() {
        // Small amount, should not reach target
        let price_current = tick_to_sqrt_price_x64(1000).unwrap();
        let price_target = tick_to_sqrt_price_x64(0).unwrap();
        let liquidity = 10_000_000_000_000u128;

        let result =
            compute_swap_step(price_current, price_target, liquidity, 100, 30, true).unwrap();

        assert!(result.sqrt_price_next > price_target, "should not reach target");
        assert!(result.sqrt_price_next < price_current);
        // Fee + amount_in should equal amount_remaining
        assert_eq!(
            result.fee_amount + result.amount_in,
            100,
            "fee + amount_in should equal amount_remaining"
        );
    }

    #[test]
    fn test_swap_step_exact_in_b_to_a() {
        let price_current = tick_to_sqrt_price_x64(0).unwrap();
        let price_target = tick_to_sqrt_price_x64(100).unwrap();
        let liquidity = 10_000_000_000u128;

        let result =
            compute_swap_step(price_current, price_target, liquidity, u64::MAX, 30, true).unwrap();

        assert_eq!(result.sqrt_price_next, price_target);
    }

    #[test]
    fn test_swap_step_exact_out_a_to_b() {
        let price_current = tick_to_sqrt_price_x64(100).unwrap();
        let price_target = tick_to_sqrt_price_x64(0).unwrap();
        let liquidity = 10_000_000_000u128;

        // Use a larger output amount so the input and fee are non-zero.
        let result =
            compute_swap_step(price_current, price_target, liquidity, 1_000_000, 30, false)
                .unwrap();

        assert!(result.amount_out <= 1_000_000);
        assert!(result.amount_in > 0);
        assert!(result.fee_amount > 0);
    }

    #[test]
    fn test_swap_step_exact_out_b_to_a() {
        let price_current = tick_to_sqrt_price_x64(0).unwrap();
        let price_target = tick_to_sqrt_price_x64(100).unwrap();
        let liquidity = 10_000_000_000u128;

        let result =
            compute_swap_step(price_current, price_target, liquidity, 100, 30, false).unwrap();

        assert!(result.amount_out <= 100);
    }

    #[test]
    fn test_swap_step_price_never_exceeds_target_a_to_b() {
        let price_current = tick_to_sqrt_price_x64(500).unwrap();
        let price_target = tick_to_sqrt_price_x64(-500).unwrap();
        let liquidity = 1_000_000_000u128;

        for amount in [1, 100, 10_000, 1_000_000, u64::MAX] {
            let result =
                compute_swap_step(price_current, price_target, liquidity, amount, 30, true)
                    .unwrap();
            assert!(
                result.sqrt_price_next >= price_target,
                "price went below target for amount={amount}"
            );
        }
    }

    #[test]
    fn test_swap_step_price_never_exceeds_target_b_to_a() {
        let price_current = tick_to_sqrt_price_x64(-500).unwrap();
        let price_target = tick_to_sqrt_price_x64(500).unwrap();
        let liquidity = 1_000_000_000u128;

        for amount in [1, 100, 10_000, 1_000_000, u64::MAX] {
            let result =
                compute_swap_step(price_current, price_target, liquidity, amount, 30, true)
                    .unwrap();
            assert!(
                result.sqrt_price_next <= price_target,
                "price went above target for amount={amount}"
            );
        }
    }

    #[test]
    fn test_swap_step_fee_zero() {
        let price_current = tick_to_sqrt_price_x64(100).unwrap();
        let price_target = tick_to_sqrt_price_x64(0).unwrap();
        let liquidity = 10_000_000_000u128;

        let result =
            compute_swap_step(price_current, price_target, liquidity, 1_000_000, 0, true).unwrap();

        assert_eq!(result.fee_amount, 0);
    }

    #[test]
    fn test_swap_step_exact_in_full_consume_fee_identity() {
        // When we don't reach target: fee + amount_in == amount_remaining
        let price_current = tick_to_sqrt_price_x64(10000).unwrap();
        let price_target = tick_to_sqrt_price_x64(0).unwrap();
        let liquidity = 1_000_000_000_000_000u128;
        let amount = 50_000u64;

        let result =
            compute_swap_step(price_current, price_target, liquidity, amount, 30, true).unwrap();

        if result.sqrt_price_next != price_target {
            assert_eq!(
                result.fee_amount + result.amount_in,
                amount,
                "partial fill: fee + in should equal remaining"
            );
        }
    }
}