riptide-amm-math 2.0.1

use super::SingleSideLiquidity;

use super::error::{CoreError, ARITHMETIC_OVERFLOW, INVALID_ORACLE_DATA};

use super::guards::{check_guards, check_oracle_validity, GuardParams};

use super::oracle::{build_liquidity, build_price, consume_liquidity, OraclePayload};

use borsh::BorshDeserialize;

use riptide_amm_macros::alias;

#[cfg(feature = "wasm")]
use riptide_amm_macros::wasm_expose;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub enum QuoteType {
    TokenAExactIn,
    TokenAExactOut,
    TokenBExactIn,
    TokenBExactOut,
}

impl QuoteType {
    pub(crate) fn new(amount_is_token_a: bool, amount_is_input: bool) -> Self {
        match (amount_is_token_a, amount_is_input) {
            (true, true) => QuoteType::TokenAExactIn,
            (true, false) => QuoteType::TokenAExactOut,
            (false, true) => QuoteType::TokenBExactIn,
            (false, false) => QuoteType::TokenBExactOut,
        }
    }

    pub fn exact_in(&self) -> bool {
        matches!(self, QuoteType::TokenAExactIn | QuoteType::TokenBExactIn)
    }

    pub fn exact_out(&self) -> bool {
        matches!(self, QuoteType::TokenAExactOut | QuoteType::TokenBExactOut)
    }

    #[alias(output_is_token_b, a_to_b)]
    pub fn input_is_token_a(&self) -> bool {
        matches!(self, QuoteType::TokenAExactIn | QuoteType::TokenBExactOut)
    }

    #[alias(output_is_token_a, b_to_a)]
    pub fn input_is_token_b(&self) -> bool {
        matches!(self, QuoteType::TokenBExactIn | QuoteType::TokenAExactOut)
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub struct Quote {
    pub amount_in: u64,
    pub amount_out: u64,
    pub quote_type: QuoteType,
}

impl Quote {
    pub fn apply_to_reserves(
        &self,
        reserves_a: u64,
        reserves_b: u64,
    ) -> Result<(u64, u64), CoreError> {
        if self.quote_type.input_is_token_a() {
            let post_a = reserves_a
                .checked_add(self.amount_in)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            let post_b = reserves_b
                .checked_sub(self.amount_out)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            Ok((post_a, post_b))
        } else {
            let post_a = reserves_a
                .checked_sub(self.amount_out)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            let post_b = reserves_b
                .checked_add(self.amount_in)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            Ok((post_a, post_b))
        }
    }
}

#[derive(Default, Debug, Clone, Eq, PartialEq)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub struct Prices {
    pub oracle_price_q64_64: u128,
    pub best_bid_price_q64_64: u128,
    pub best_ask_price_q64_64: u128,
    pub ask_spread_per_m: i32,
    pub bid_spread_per_m: i32,
}

#[derive(Default, Debug, Clone, Eq, PartialEq)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub struct Price {
    pub oracle_price_q64_64: u128,
    pub best_price_q64_64: u128,
    pub spread_per_m: i32,
}

#[derive(Debug, Clone, Eq, PartialEq)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub struct GuardedQuote {
    pub quote: Quote,
    pub post_reserves_a: u64,
    pub post_reserves_b: u64,
    pub price: Price,
}

pub type QuoteError = &'static str;

#[derive(Debug)]
struct InnerQuoteResult {
    quote: Quote,
    payload: OraclePayload,
    liquidity: SingleSideLiquidity,
}

fn quote(
    amount: u64,
    amount_is_token_a: bool,
    amount_is_input: bool,
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
) -> Result<InnerQuoteResult, CoreError> {
    let mut oracle_data = oracle_data;
    let payload = OraclePayload::deserialize(&mut oracle_data).map_err(|_| INVALID_ORACLE_DATA)?;

    let quote_type = QuoteType::new(amount_is_token_a, amount_is_input);

    let liquidity = build_liquidity(
        &payload,
        quote_type,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
    )?;

    let quote = consume_liquidity(amount, quote_type, &liquidity)?;

    Ok(InnerQuoteResult {
        quote,
        payload,
        liquidity,
    })
}

#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn quote_exact_in(
    amount: u64,
    amount_is_token_a: bool,
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
) -> Result<Quote, CoreError> {
    quote(
        amount,
        amount_is_token_a,
        true,
        oracle_data,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
    )
    .map(|r| r.quote)
}

#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn quote_exact_out(
    amount: u64,
    amount_is_token_a: bool,
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
) -> Result<Quote, CoreError> {
    quote(
        amount,
        amount_is_token_a,
        false,
        oracle_data,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
    )
    .map(|r| r.quote)
}

fn quote_with_guards(
    amount: u64,
    amount_is_token_a: bool,
    amount_is_input: bool,
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
    current_slot: u64,
    params: &GuardParams,
) -> Result<GuardedQuote, QuoteError> {
    check_oracle_validity(current_slot, params.valid_until)?;

    let inner = quote(
        amount,
        amount_is_token_a,
        amount_is_input,
        oracle_data,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
    )?;

    let quote_type = QuoteType::new(amount_is_token_a, amount_is_input);

    let price = build_price(
        &inner.liquidity,
        &inner.payload.data,
        quote_type,
        reserves_a,
        reserves_b,
    )?;

    let (post_reserves_a, post_reserves_b) =
        inner.quote.apply_to_reserves(reserves_a, reserves_b)?;

    check_guards(post_reserves_a, post_reserves_b, &price, params)?;

    Ok(GuardedQuote {
        quote: inner.quote,
        post_reserves_a,
        post_reserves_b,
        price,
    })
}

pub fn quote_exact_in_with_guards(
    amount: u64,
    amount_is_token_a: bool,
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
    current_slot: u64,
    params: &GuardParams,
) -> Result<GuardedQuote, QuoteError> {
    quote_with_guards(
        amount,
        amount_is_token_a,
        true,
        oracle_data,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
        current_slot,
        params,
    )
}

pub fn quote_exact_out_with_guards(
    amount: u64,
    amount_is_token_a: bool,
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
    current_slot: u64,
    params: &GuardParams,
) -> Result<GuardedQuote, QuoteError> {
    quote_with_guards(
        amount,
        amount_is_token_a,
        false,
        oracle_data,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
        current_slot,
        params,
    )
}

#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn bid_price(
    oracle_data: &[u8],
    liquidity: SingleSideLiquidity,
    reserves_a: u64,
    reserves_b: u64,
) -> Result<Price, CoreError> {
    let mut oracle_data = oracle_data;
    let payload = OraclePayload::deserialize(&mut oracle_data).map_err(|_| INVALID_ORACLE_DATA)?;
    build_price(
        &liquidity,
        &payload.data,
        QuoteType::TokenAExactIn,
        reserves_a,
        reserves_b,
    )
}

#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn ask_price(
    oracle_data: &[u8],
    liquidity: SingleSideLiquidity,
    reserves_a: u64,
    reserves_b: u64,
) -> Result<Price, CoreError> {
    let mut oracle_data = oracle_data;
    let payload = OraclePayload::deserialize(&mut oracle_data).map_err(|_| INVALID_ORACLE_DATA)?;
    build_price(
        &liquidity,
        &payload.data,
        QuoteType::TokenBExactIn,
        reserves_a,
        reserves_b,
    )
}

#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn bid_liquidity(
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
) -> Result<SingleSideLiquidity, CoreError> {
    let mut oracle_data = oracle_data;
    let payload = OraclePayload::deserialize(&mut oracle_data).map_err(|_| INVALID_ORACLE_DATA)?;
    build_liquidity(
        &payload,
        QuoteType::TokenAExactIn,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
    )
}

#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn ask_liquidity(
    oracle_data: &[u8],
    reserves_a: u64,
    reserves_b: u64,
    skew_cliff_min_per_m: i32,
    skew_cliff_max_per_m: i32,
) -> Result<SingleSideLiquidity, CoreError> {
    let mut oracle_data = oracle_data;
    let payload = OraclePayload::deserialize(&mut oracle_data).map_err(|_| INVALID_ORACLE_DATA)?;
    build_liquidity(
        &payload,
        QuoteType::TokenBExactIn,
        reserves_a,
        reserves_b,
        skew_cliff_min_per_m,
        skew_cliff_max_per_m,
    )
}

#[cfg(test)]
mod tests {
    use super::{
        super::guards::{
            INVENTORY_IMBALANCE, ORACLE_EXPIRED, ORACLE_PRICE_BELOW_MIN, SPREAD_BELOW_MIN,
        },
        *,
    };
    use borsh::BorshSerialize;
    use rstest::rstest;

    use super::super::oracle::{
        OracleData, SkewMode, ORACLE_DATA_LEN, ORACLE_PAYLOAD_LEN, SKEW_OFFSET,
    };

    fn flat_oracle_data(price_q64_64: u128) -> [u8; ORACLE_PAYLOAD_LEN] {
        let data = OracleData::FlatPrice { price_q64_64 };
        let skew = SkewMode::None;

        let mut buf = [0u8; ORACLE_PAYLOAD_LEN];

        let mut data_slice = &mut buf[..ORACLE_DATA_LEN];
        data.serialize(&mut data_slice).unwrap();

        let mut skew_slice = &mut buf[SKEW_OFFSET..];
        skew.serialize(&mut skew_slice).unwrap();

        buf
    }

    fn pass_through_params() -> GuardParams {
        GuardParams {
            max_inventory_imbalance_per_m: i32::MAX,
            max_a_inventory_per_m: 0,
            max_b_inventory_per_m: 0,
            min_spread_per_m: i32::MIN,
            min_oracle_price: 0,
            max_oracle_price: u128::MAX,
            valid_until: u64::MAX,
        }
    }

    #[rstest]
    #[case::a_to_b_ok(100, 80, QuoteType::TokenAExactIn, 1000, 1000, Ok((1100, 920)))]
    #[case::b_to_a_ok(80, 100, QuoteType::TokenBExactIn, 1000, 1000, Ok((900, 1080)))]
    #[case::a_to_b_exact_out(100, 80, QuoteType::TokenBExactOut, 1000, 1000, Ok((1100, 920)))]
    #[case::b_to_a_exact_out(80, 100, QuoteType::TokenAExactOut, 1000, 1000, Ok((900, 1080)))]
    #[case::output_underflow_a(
        100,
        2000,
        QuoteType::TokenAExactIn,
        1000,
        1000,
        Err(ARITHMETIC_OVERFLOW)
    )]
    #[case::output_underflow_b(
        100,
        2000,
        QuoteType::TokenBExactIn,
        1000,
        1000,
        Err(ARITHMETIC_OVERFLOW)
    )]
    #[case::input_overflow_a(
        1,
        0,
        QuoteType::TokenAExactIn,
        u64::MAX,
        1000,
        Err(ARITHMETIC_OVERFLOW)
    )]
    #[case::input_overflow_b(
        1,
        0,
        QuoteType::TokenBExactIn,
        1000,
        u64::MAX,
        Err(ARITHMETIC_OVERFLOW)
    )]
    fn test_apply_to_reserves(
        #[case] amount_in: u64,
        #[case] amount_out: u64,
        #[case] quote_type: QuoteType,
        #[case] reserves_a: u64,
        #[case] reserves_b: u64,
        #[case] expected: Result<(u64, u64), CoreError>,
    ) {
        let quote = Quote {
            amount_in,
            amount_out,
            quote_type,
        };

        let result = quote.apply_to_reserves(reserves_a, reserves_b);

        assert_eq!(result, expected);
    }

    #[test]
    fn test_guards_happy_path() {
        let oracle = flat_oracle_data(1 << 64);
        let params = pass_through_params();

        let result = quote_exact_in_with_guards(100, true, &oracle, 1000, 1000, 0, 0, 0, &params);

        let guarded = result.unwrap();
        assert_eq!(guarded.quote.amount_in, 100);
        assert_eq!(guarded.quote.amount_out, 100);
        assert_eq!(guarded.post_reserves_a, 1100);
        assert_eq!(guarded.post_reserves_b, 900);
    }

    #[test]
    fn test_guards_oracle_expired() {
        let oracle = flat_oracle_data(1 << 64);
        let params = GuardParams {
            valid_until: 5,
            ..pass_through_params()
        };

        let result = quote_exact_in_with_guards(100, true, &oracle, 1000, 1000, 0, 0, 10, &params);

        assert_eq!(result, Err(ORACLE_EXPIRED));
    }

    #[test]
    fn test_guards_inventory_guard_fail() {
        let oracle = flat_oracle_data(1 << 64);
        let params = GuardParams {
            max_inventory_imbalance_per_m: 10_000,
            ..pass_through_params()
        };

        let result = quote_exact_in_with_guards(100, true, &oracle, 1500, 500, 0, 0, 0, &params);

        assert_eq!(result, Err(INVENTORY_IMBALANCE));
    }

    #[test]
    fn test_guards_spread_guard_fail() {
        let oracle = flat_oracle_data(1 << 64);
        let params = GuardParams {
            min_spread_per_m: 100,
            ..pass_through_params()
        };

        let result = quote_exact_in_with_guards(100, true, &oracle, 1000, 1000, 0, 0, 0, &params);

        assert_eq!(result, Err(SPREAD_BELOW_MIN));
    }

    #[test]
    fn test_guards_price_below_min() {
        let price = 1u128 << 64;
        let oracle = flat_oracle_data(price);
        let params = GuardParams {
            min_oracle_price: price + 1,
            ..pass_through_params()
        };

        let result = quote_exact_in_with_guards(100, true, &oracle, 1000, 1000, 0, 0, 0, &params);

        assert_eq!(result, Err(ORACLE_PRICE_BELOW_MIN));
    }

    #[test]
    fn test_guards_invalid_oracle() {
        let params = pass_through_params();

        let result = quote_exact_in_with_guards(100, true, &[0u8; 4], 1000, 1000, 0, 0, 0, &params);

        assert_eq!(result, Err(INVALID_ORACLE_DATA));
    }
}