riptide_amm/math/oracle/

mod.rs

use std::cmp::min;

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

use borsh::{BorshDeserialize, BorshSerialize};

use super::{
    error::{CoreError, ARITHMETIC_OVERFLOW},
    quote::{Quote, QuoteType},
    token::{a_to_b, b_to_a},
};

mod book;
mod flat;
mod spread;

use book::{book_liquidity, BookSpacingType};
use flat::flat_liquidity;
use spread::spread_liquidity;

// Represents the liquidity on one side of the order book (price and amount of out_token).
// This is explicitly a fixed-size array to avoid heap allocation (which Vec does).
// The `len` field is the actual number of initialized elements in the array.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub(crate) struct SingleSideLiquidity {
    items: [(u128, u64); 32],
    len: usize,
}

#[allow(dead_code)]
impl SingleSideLiquidity {
    pub fn new() -> Self {
        Self {
            items: [(0, 0); 32],
            len: 0,
        }
    }

    pub fn push(&mut self, item: (u128, u64)) -> bool {
        if self.len < 32 {
            self.items[self.len] = item;
            self.len += 1;
            true
        } else {
            false
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn as_slice(&self) -> &[(u128, u64)] {
        &self.items[..self.len]
    }

    pub fn to_vec(&self) -> Vec<(u128, u64)> {
        self.as_slice().to_vec()
    }
}

#[cfg_attr(feature = "wasm", wasm_expose)]
pub const PER_1M_DENOMINATOR: i32 = 1_000_000;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[cfg_attr(true, derive(BorshDeserialize, BorshSerialize))]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub enum OracleData {
    Empty,
    Flat {
        price_q64_64: u128,
    },
    Spread {
        price_q64_64: u128,
        spread_a_to_b_per_1m: i32,
        spread_b_to_a_per_1m: i32,
    },
    Book {
        price_q64_64: u128,
        // The spacing between each level of liquidity
        spacing: BookSpacingType,
        // Expressed as a fraction of the total reserves
        bid_liquidity_per_1m: [u32; 32],
        ask_liquidity_per_1m: [u32; 32],
    },
}

impl OracleData {
    pub(crate) fn swap(
        &self,
        amount: u64,
        amount_is_token_a: bool,
        amount_is_input: bool,
        reserves_a: u64,
        reserves_b: u64,
    ) -> Result<Quote, CoreError> {
        let liquidity: SingleSideLiquidity = match self {
            OracleData::Empty => SingleSideLiquidity::new(),
            OracleData::Flat { price_q64_64 } => flat_liquidity(
                *price_q64_64,
                amount_is_token_a,
                amount_is_input,
                reserves_a,
                reserves_b,
            )?,
            OracleData::Spread {
                price_q64_64,
                spread_a_to_b_per_1m,
                spread_b_to_a_per_1m,
            } => spread_liquidity(
                *price_q64_64,
                *spread_a_to_b_per_1m,
                *spread_b_to_a_per_1m,
                amount_is_token_a,
                amount_is_input,
                reserves_a,
                reserves_b,
            )?,
            OracleData::Book {
                price_q64_64,
                spacing,
                bid_liquidity_per_1m,
                ask_liquidity_per_1m,
            } => book_liquidity(
                amount_is_token_a,
                amount_is_input,
                *price_q64_64,
                *spacing,
                bid_liquidity_per_1m,
                ask_liquidity_per_1m,
                reserves_a,
                reserves_b,
            )?,
        };

        consume_liquidity(
            amount,
            amount_is_token_a,
            amount_is_input,
            &liquidity.as_slice(),
        )
    }
}

pub(crate) fn consume_liquidity(
    amount: u64,
    amount_is_token_a: bool,
    amount_is_input: bool,
    liquidity: &[(u128, u64)],
) -> Result<Quote, CoreError> {
    let mut remaining_amount = amount;
    let mut other_amount: u64 = 0;

    // Start consuming the liquidity from the best price (first)
    // Stop when either the liquidity runs out or the full amount is consumed
    // `liquidity` represents the amount of output token that is available at the given price

    for (price, liquidity) in liquidity {
        let (step_specified_amount, step_other_amount) = if amount_is_input {
            // ExactIn
            let max_amount = if amount_is_token_a {
                b_to_a(*liquidity, (*price).into(), true)?
            } else {
                a_to_b(*liquidity, (*price).into(), true)?
            };

            let step_specified_amount = min(remaining_amount, max_amount);
            let step_other_amount = if amount_is_token_a {
                a_to_b(step_specified_amount, (*price).into(), false)?
            } else {
                b_to_a(step_specified_amount, (*price).into(), false)?
            };

            (step_specified_amount, min(step_other_amount, *liquidity))
        } else {
            // ExactOut

            let step_specified_amount = min(remaining_amount, *liquidity);
            let step_other_amount = if amount_is_token_a {
                a_to_b(step_specified_amount, (*price).into(), true)?
            } else {
                b_to_a(step_specified_amount, (*price).into(), true)?
            };

            (step_specified_amount, step_other_amount)
        };

        remaining_amount = remaining_amount
            .checked_sub(step_specified_amount)
            .ok_or(ARITHMETIC_OVERFLOW)?;
        other_amount = other_amount
            .checked_add(step_other_amount)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        if remaining_amount == 0 {
            break;
        }
    }

    let consumed_amount = amount - remaining_amount;

    if amount_is_input {
        Ok(Quote {
            amount_in: consumed_amount,
            amount_out: other_amount,
            quote_type: QuoteType::ExactIn,
        })
    } else {
        Ok(Quote {
            amount_in: other_amount,
            amount_out: consumed_amount,
            quote_type: QuoteType::ExactOut,
        })
    }
}

#[cfg(all(test, feature = "lib"))]
mod tests {
    use super::*;
    use rstest::rstest;

    #[rstest]
    fn test_empty(
        #[values(true, false)] amount_is_token_a: bool,
        #[values(true, false)] amount_is_input: bool,
    ) {
        let quote = OracleData::Empty.swap(100, amount_is_token_a, amount_is_input, 1000, 1000);
        let quote_type = if amount_is_input {
            QuoteType::ExactIn
        } else {
            QuoteType::ExactOut
        };
        assert_eq!(
            quote,
            Ok(Quote {
                amount_in: 0,
                amount_out: 0,
                quote_type,
            })
        );
    }

    // Selling exact token_a. The lower the the price, the less token_b you get.
    #[rstest]
    #[case(100, Ok(Quote { amount_in: 100, amount_out: 100, quote_type: QuoteType::ExactIn }))]
    #[case(500, Ok(Quote { amount_in: 500, amount_out: 300, quote_type: QuoteType::ExactIn }))]
    #[case(1100, Ok(Quote { amount_in: 1100, amount_out: 600, quote_type: QuoteType::ExactIn }))]
    #[case(2500, Ok(Quote { amount_in: 2500, amount_out: 950, quote_type: QuoteType::ExactIn }))]
    #[case(10000, Ok(Quote { amount_in: 5100, amount_out: 1600, quote_type: QuoteType::ExactIn }))]
    fn test_consume_liquidity_input_token_a(
        #[case] amount: u64,
        #[case] expected: Result<Quote, CoreError>,
    ) {
        let liquidity = vec![(1 << 64, 100), ((1 << 64) / 2, 500), ((1 << 64) / 4, 1000)];
        let quote = consume_liquidity(amount, true, true, &liquidity);
        assert_eq!(quote, expected);
    }

    // Selling exact token_b. The higher the the price, the less token_a you get.
    #[rstest]
    #[case(100, Ok(Quote { amount_in: 100, amount_out: 100, quote_type: QuoteType::ExactIn }))]
    #[case(500, Ok(Quote { amount_in: 500, amount_out: 300, quote_type: QuoteType::ExactIn }))]
    #[case(1100, Ok(Quote { amount_in: 1100, amount_out: 600, quote_type: QuoteType::ExactIn }))]
    #[case(2500, Ok(Quote { amount_in: 2500, amount_out: 950, quote_type: QuoteType::ExactIn }))]
    #[case(10000, Ok(Quote { amount_in: 5100, amount_out: 1600, quote_type: QuoteType::ExactIn }))]
    fn test_consume_liquidity_input_token_b(
        #[case] amount: u64,
        #[case] expected: Result<Quote, CoreError>,
    ) {
        let liquidity = vec![(1 << 64, 100), (2 << 64, 500), (4 << 64, 1000)];
        let quote = consume_liquidity(amount, false, true, &liquidity);
        assert_eq!(quote, expected);
    }

    // Buying exact token_a. The higher the the price, the more token_b you pay.
    #[rstest]
    #[case(100, Ok(Quote { amount_in: 100, amount_out: 100, quote_type: QuoteType::ExactOut }))]
    #[case(300, Ok(Quote { amount_in: 500, amount_out: 300, quote_type: QuoteType::ExactOut }))]
    #[case(600, Ok(Quote { amount_in: 1100, amount_out: 600, quote_type: QuoteType::ExactOut }))]
    #[case(950, Ok(Quote { amount_in: 2500, amount_out: 950, quote_type: QuoteType::ExactOut }))]
    #[case(2000, Ok(Quote { amount_in: 5100, amount_out: 1600, quote_type: QuoteType::ExactOut }))]
    fn test_consume_liquidity_output_token_a(
        #[case] amount: u64,
        #[case] expected: Result<Quote, CoreError>,
    ) {
        let liquidity = vec![(1 << 64, 100), (2 << 64, 500), (4 << 64, 1000)];
        let quote = consume_liquidity(amount, true, false, &liquidity);
        assert_eq!(quote, expected);
    }

    // Buying exact token_b. The lower the the price, the more token_a you pay.
    #[rstest]
    #[case(100, Ok(Quote { amount_in: 100, amount_out: 100, quote_type: QuoteType::ExactOut }))]
    #[case(300, Ok(Quote { amount_in: 500, amount_out: 300, quote_type: QuoteType::ExactOut }))]
    #[case(600, Ok(Quote { amount_in: 1100, amount_out: 600, quote_type: QuoteType::ExactOut }))]
    #[case(950, Ok(Quote { amount_in: 2500, amount_out: 950, quote_type: QuoteType::ExactOut }))]
    #[case(2000, Ok(Quote { amount_in: 5100, amount_out: 1600, quote_type: QuoteType::ExactOut }))]
    fn test_consume_liquidity_output_token_b(
        #[case] amount: u64,
        #[case] expected: Result<Quote, CoreError>,
    ) {
        let liquidity = vec![(1 << 64, 100), ((1 << 64) / 2, 500), ((1 << 64) / 4, 1000)];
        let quote = consume_liquidity(amount, false, false, &liquidity);
        assert_eq!(quote, expected);
    }

    #[rstest]
    #[case((1 << 64) / 8, true, true, Ok(Quote { amount_in: 100, amount_out: 12, quote_type: QuoteType::ExactIn }))]
    #[case((1 << 64) / 8, true, false, Ok(Quote { amount_in: 13, amount_out: 100, quote_type: QuoteType::ExactOut }))]
    #[case(8 << 64, false, true, Ok(Quote { amount_in: 100, amount_out: 12, quote_type: QuoteType::ExactIn }))]
    #[case(8 << 64, false, false, Ok(Quote { amount_in: 13, amount_out: 100, quote_type: QuoteType::ExactOut }))]
    fn test_consume_liquidity_rounding_direction(
        #[case] price: u128,
        #[case] amount_is_token_a: bool,
        #[case] amount_is_input: bool,
        #[case] expected: Result<Quote, CoreError>,
    ) {
        let liquidity = vec![(price, 1000)];
        let result = consume_liquidity(100, amount_is_token_a, amount_is_input, &liquidity);
        assert_eq!(result, expected);
    }
}