riptide-amm-math 1.0.3

use borsh::{BorshDeserialize, BorshSerialize};
#[cfg(feature = "wasm")]
use riptide_amm_macros::wasm_expose;

use super::{
    super::{
        error::{CoreError, AMOUNT_EXCEEDS_MAX_I32, ARITHMETIC_OVERFLOW},
        quote::QuoteType,
    },
    SingleSideLiquidity, PER_M_DENOMINATOR,
};
use ethnum::U256;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[cfg_attr(true, derive(BorshDeserialize, BorshSerialize))]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub enum SkewExponent {
    Linear,
    Quadratic,
    Cubic,
}

impl SkewExponent {
    pub fn value(&self) -> u32 {
        match self {
            Self::Linear => 1,
            Self::Quadratic => 2,
            Self::Cubic => 3,
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[cfg_attr(true, derive(BorshDeserialize, BorshSerialize))]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub enum SkewMode {
    None,
    Polynomial {
        exponent: SkewExponent,
        positive_bid_per_m: u32,
        negative_bid_per_m: u32,
        positive_ask_per_m: u32,
        negative_ask_per_m: u32,
    },
}

fn select_intensity(
    deviation_per_m: i32,
    a_to_b: bool,
    positive_bid_per_m: u32,
    negative_bid_per_m: u32,
    positive_ask_per_m: u32,
    negative_ask_per_m: u32,
) -> u32 {
    match (deviation_per_m >= 0, a_to_b) {
        (true, true) => positive_bid_per_m,
        (false, true) => negative_bid_per_m,
        (true, false) => positive_ask_per_m,
        (false, false) => negative_ask_per_m,
    }
}

impl SkewMode {
    /// Computes the skew value in per_m units.
    /// `deviation_per_m` ranges from -1_000_000 to +1_000_000.
    /// Returns the skew adjustment to apply: bid_spread += skew, ask_spread -= skew.
    pub(crate) fn compute_skew_per_m(
        &self,
        deviation_per_m: i32,
        quote_type: QuoteType,
    ) -> Result<i32, CoreError> {
        match self {
            Self::None => Ok(0),
            Self::Polynomial {
                exponent,
                positive_bid_per_m,
                negative_bid_per_m,
                positive_ask_per_m,
                negative_ask_per_m,
            } => {
                let a_to_b = quote_type.a_to_b();
                let intensity = select_intensity(
                    deviation_per_m,
                    a_to_b,
                    *positive_bid_per_m,
                    *negative_bid_per_m,
                    *positive_ask_per_m,
                    *negative_ask_per_m,
                );
                let sign = deviation_per_m.signum();
                let abs_dev = deviation_per_m.unsigned_abs() as u128;
                let exp = exponent.value();

                let numerator = abs_dev
                    .checked_pow(exp)
                    .ok_or(ARITHMETIC_OVERFLOW)?
                    .checked_mul(intensity as u128)
                    .ok_or(ARITHMETIC_OVERFLOW)?;
                let denominator = (PER_M_DENOMINATOR as u128)
                    .checked_pow(exp)
                    .ok_or(ARITHMETIC_OVERFLOW)?;

                let quotient = numerator
                    .checked_div(denominator)
                    .ok_or(ARITHMETIC_OVERFLOW)?;
                let remainder = numerator
                    .checked_rem(denominator)
                    .ok_or(ARITHMETIC_OVERFLOW)?;
                let abs_result = if remainder > 0 {
                    quotient.checked_add(1).ok_or(ARITHMETIC_OVERFLOW)?
                } else {
                    quotient
                };

                let result = i32::try_from(abs_result).map_err(|_| AMOUNT_EXCEEDS_MAX_I32)?;
                sign.checked_mul(result).ok_or(ARITHMETIC_OVERFLOW)
            }
        }
    }
}

/// Applies skew to all prices in SingleSideLiquidity.
/// adjusted_price = price * (PER_M - skew_per_m) / PER_M
pub(crate) fn apply_skew_to_liquidity(
    liquidity: SingleSideLiquidity,
    skew_per_m: i32,
    quote_type: QuoteType,
) -> Result<SingleSideLiquidity, CoreError> {
    let clamped = skew_per_m.clamp(-PER_M_DENOMINATOR, PER_M_DENOMINATOR);
    if clamped == 0 {
        return Ok(liquidity);
    }
    let a_to_b = quote_type.a_to_b();
    let widening = a_to_b == (clamped > 0);
    let abs_skew = U256::from(clamped.unsigned_abs());
    let denom = U256::from(PER_M_DENOMINATOR as u64);
    let mut result = SingleSideLiquidity::new();
    for &(price, amount) in liquidity.as_slice() {
        let numerator = U256::from(price)
            .checked_mul(abs_skew)
            .ok_or(ARITHMETIC_OVERFLOW)?;
        let quotient = numerator.checked_div(denom).ok_or(ARITHMETIC_OVERFLOW)?;
        let remainder = numerator.checked_rem(denom).ok_or(ARITHMETIC_OVERFLOW)?;
        let delta: u128 = if widening && remainder > U256::ZERO {
            quotient.checked_add(U256::ONE).ok_or(ARITHMETIC_OVERFLOW)?
        } else {
            quotient
        }
        .try_into()
        .map_err(|_| ARITHMETIC_OVERFLOW)?;
        let adjusted_price = if clamped > 0 {
            price.checked_sub(delta).ok_or(ARITHMETIC_OVERFLOW)?
        } else {
            price.checked_add(delta).ok_or(ARITHMETIC_OVERFLOW)?
        };
        result.push((adjusted_price, amount));
    }
    Ok(result)
}

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

    const PRICE_ONE: u128 = 1 << 64;

    #[rstest]
    #[case(SkewMode::None, 200_000, QuoteType::TokenAExactIn, 0)]
    // symmetric Linear
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 500_000, negative_bid_per_m: 500_000, positive_ask_per_m: 500_000, negative_ask_per_m: 500_000 }, 200_000, QuoteType::TokenAExactIn, 100_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 500_000, negative_bid_per_m: 500_000, positive_ask_per_m: 500_000, negative_ask_per_m: 500_000 }, -200_000, QuoteType::TokenAExactIn, -100_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, 1_000_000, QuoteType::TokenAExactIn, 1_000_000)]
    // symmetric Quadratic: 200_000^2 * 1_000_000 / 1_000_000^2 = 40_000
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Quadratic, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, 200_000, QuoteType::TokenAExactIn, 40_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Quadratic, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, -200_000, QuoteType::TokenAExactIn, -40_000)]
    // symmetric Cubic: 200_000^3 * 1_000_000 / 1_000_000^3 = 8_000
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Cubic, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, 200_000, QuoteType::TokenAExactIn, 8_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Cubic, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, -200_000, QuoteType::TokenAExactIn, -8_000)]
    // Cubic at max deviation: 1_000_000^3 * 1_000_000 / 1_000_000^3 = 1_000_000
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Cubic, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, 1_000_000, QuoteType::TokenAExactIn, 1_000_000)]
    // no remainder, no rounding
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, 1, QuoteType::TokenAExactIn, 1)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 1_000_000, negative_bid_per_m: 1_000_000, positive_ask_per_m: 1_000_000, negative_ask_per_m: 1_000_000 }, -1, QuoteType::TokenAExactOut, -1)]
    // 700_001 * 300_000 / 1_000_000 = 210_000.3
    // round up (a_to_b): -> 210_001
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 300_000, negative_bid_per_m: 300_000, positive_ask_per_m: 300_000, negative_ask_per_m: 300_000 }, 700_001, QuoteType::TokenAExactIn, 210_001)]
    // always round away from zero: -> 210_001
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 300_000, negative_bid_per_m: 300_000, positive_ask_per_m: 300_000, negative_ask_per_m: 300_000 }, 700_001, QuoteType::TokenAExactOut, 210_001)]
    // negative deviation: round up -> -210_001
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 300_000, negative_bid_per_m: 300_000, positive_ask_per_m: 300_000, negative_ask_per_m: 300_000 }, -700_001, QuoteType::TokenAExactIn, -210_001)]
    // negative deviation: always round away from zero -> -210_001
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 300_000, negative_bid_per_m: 300_000, positive_ask_per_m: 300_000, negative_ask_per_m: 300_000 }, -700_001, QuoteType::TokenAExactOut, -210_001)]
    // asymmetric Linear: 4 quadrants
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 500_000, QuoteType::TokenAExactIn, 50_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, -500_000, QuoteType::TokenAExactIn, -100_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 500_000, QuoteType::TokenBExactIn, 150_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, -500_000, QuoteType::TokenBExactIn, -200_000)]
    // zero deviation
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 0, QuoteType::TokenAExactIn, 0)]
    // rounding with asymmetric: bid ceil 700_001 * 100_000 / 1_000_000 = 70_000.1 -> 70_001
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 700_001, QuoteType::TokenAExactIn, 70_001)]
    // rounding with asymmetric: always away from zero 700_001 * 300_000 / 1_000_000 = 210_000.3 ->
    // 210_001
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 700_001, QuoteType::TokenBExactIn, 210_001)]
    // one quadrant zero: positive_bid=0 -> skew=0
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Linear, positive_bid_per_m: 0, negative_bid_per_m: 500_000, positive_ask_per_m: 500_000, negative_ask_per_m: 500_000 }, 500_000, QuoteType::TokenAExactIn, 0)]
    // asymmetric Quadratic: 500_000^2 * selected_intensity / 1_000_000^2
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Quadratic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 500_000, QuoteType::TokenAExactIn, 25_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Quadratic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, -500_000, QuoteType::TokenAExactIn, -50_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Quadratic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 500_000, QuoteType::TokenBExactIn, 75_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Quadratic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, -500_000, QuoteType::TokenBExactIn, -100_000)]
    // asymmetric Cubic: 500_000^3 * selected_intensity / 1_000_000^3
    // 500_000^3 = 125_000_000_000_000_000, * 100_000 / 10^18 = 12_500
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Cubic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 500_000, QuoteType::TokenAExactIn, 12_500)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Cubic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, -500_000, QuoteType::TokenAExactIn, -25_000)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Cubic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, 500_000, QuoteType::TokenBExactIn, 37_500)]
    #[case(SkewMode::Polynomial { exponent: SkewExponent::Cubic, positive_bid_per_m: 100_000, negative_bid_per_m: 200_000, positive_ask_per_m: 300_000, negative_ask_per_m: 400_000 }, -500_000, QuoteType::TokenBExactIn, -50_000)]
    fn test_compute_skew_per_m(
        #[case] skew_mode: SkewMode,
        #[case] deviation_per_m: i32,
        #[case] quote_type: QuoteType,
        #[case] expected: i32,
    ) {
        assert_eq!(
            skew_mode
                .compute_skew_per_m(deviation_per_m, quote_type)
                .unwrap(),
            expected
        );
    }

    #[rstest]
    // widening: a_to_b + positive skew, or !a_to_b + negative skew -> ceil delta
    #[case(PRICE_ONE, 0, QuoteType::TokenAExactIn, PRICE_ONE)]
    #[case(PRICE_ONE, 500_000, QuoteType::TokenAExactIn, PRICE_ONE / 2)]
    #[case(PRICE_ONE, -500_000, QuoteType::TokenBExactIn, PRICE_ONE + PRICE_ONE / 2)]
    #[case(PRICE_ONE, 250_000, QuoteType::TokenAExactIn, PRICE_ONE * 3 / 4)]
    #[case(PRICE_ONE, 1_000_000, QuoteType::TokenAExactIn, 0)]
    #[case(PRICE_ONE, -1_000_000, QuoteType::TokenBExactIn, PRICE_ONE * 2)]
    #[case(PRICE_ONE, 2_000_000, QuoteType::TokenAExactIn, 0)] // clamped to 1M
    #[case(PRICE_ONE, -2_000_000, QuoteType::TokenBExactIn, PRICE_ONE * 2)] // clamped to -1M
    // widening: delta ceiling for protocol safety
    #[case(PRICE_ONE, 999_999, QuoteType::TokenAExactIn, PRICE_ONE - (PRICE_ONE * 999_999).div_ceil(1_000_000))]
    #[case(PRICE_ONE, -999_999, QuoteType::TokenBExactIn, PRICE_ONE + (PRICE_ONE * 999_999).div_ceil(1_000_000))]
    // narrowing: a_to_b + negative skew, or !a_to_b + positive skew -> floor delta
    #[case(PRICE_ONE, -999_999, QuoteType::TokenAExactIn, PRICE_ONE + PRICE_ONE * 999_999 / 1_000_000)]
    #[case(PRICE_ONE, 999_999, QuoteType::TokenBExactIn, PRICE_ONE - PRICE_ONE * 999_999 / 1_000_000)]
    fn test_apply_skew_to_liquidity(
        #[case] price: u128,
        #[case] skew_per_m: i32,
        #[case] quote_type: QuoteType,
        #[case] expected_price: u128,
    ) {
        let liquidity = SingleSideLiquidity::from_slice(&[(price, 1000)]);
        let result = apply_skew_to_liquidity(liquidity, skew_per_m, quote_type).unwrap();
        let (result_price, result_amount) = result.as_slice()[0];
        assert_eq!(result_price, expected_price);
        assert_eq!(result_amount, 1000);
    }
}