riptide-amm-math 1.0.3

use super::super::{
    error::{
        CoreError, AMOUNT_EXCEEDS_MAX_U128, AMOUNT_EXCEEDS_MAX_U32, AMOUNT_EXCEEDS_MAX_U64,
        ARITHMETIC_OVERFLOW, INVALID_ORACLE_DATA,
    },
    oracle::{OracleData, PER_M_DENOMINATOR},
    quote::{Quote, QuoteType},
};

use borsh::{BorshDeserialize, BorshSerialize};

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

use super::SingleSideLiquidity;

pub(crate) const BOOK_LIQUIDITY_LEVELS: usize = 32;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[cfg_attr(true, derive(BorshDeserialize, BorshSerialize))]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub enum BookSpacingType {
    // Price = price_q64_64 ± (X * i * price_q64_64)
    Linear(u16),
    // Price = price_q64_64 * (1 ± X)^i
    Exponential(u16),
}

impl BookSpacingType {
    pub fn price(&self, price_q64_64: u128, i: usize, a_to_b: bool) -> Result<u128, CoreError> {
        match *self {
            BookSpacingType::Linear(spacing_per_m) => {
                let product = price_q64_64
                    .checked_mul(spacing_per_m as u128)
                    .ok_or(ARITHMETIC_OVERFLOW)?
                    .checked_mul(i as u128)
                    .ok_or(ARITHMETIC_OVERFLOW)?;
                let quotient = product
                    .checked_div(PER_M_DENOMINATOR as u128)
                    .ok_or(ARITHMETIC_OVERFLOW)?;
                let remainder = product
                    .checked_rem(PER_M_DENOMINATOR as u128)
                    .ok_or(ARITHMETIC_OVERFLOW)?;
                let spread = if !a_to_b && remainder > 0 {
                    quotient + 1
                } else {
                    quotient
                };
                if a_to_b {
                    price_q64_64.checked_sub(spread).ok_or(ARITHMETIC_OVERFLOW)
                } else {
                    price_q64_64.checked_add(spread).ok_or(ARITHMETIC_OVERFLOW)
                }
            }
            BookSpacingType::Exponential(spacing_per_m) => {
                let factor = if a_to_b {
                    PER_M_DENOMINATOR - spacing_per_m as i32
                } else {
                    PER_M_DENOMINATOR + spacing_per_m as i32
                };
                let mut base = u128::from(factor as u32)
                    .checked_shl(64)
                    .ok_or(ARITHMETIC_OVERFLOW)?
                    .checked_div(PER_M_DENOMINATOR as u128)
                    .ok_or(ARITHMETIC_OVERFLOW)?;

                let mut exp = i as u32;
                let mut result: u128 = price_q64_64;

                while exp > 0 {
                    if exp & 1 == 1 {
                        result = U256::from(result)
                            .checked_mul(U256::from(base))
                            .ok_or(ARITHMETIC_OVERFLOW)?
                            .checked_shr(64)
                            .ok_or(ARITHMETIC_OVERFLOW)?
                            .try_into()
                            .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?;
                    }

                    base = U256::from(base)
                        .checked_mul(U256::from(base))
                        .ok_or(ARITHMETIC_OVERFLOW)?
                        .checked_shr(64)
                        .ok_or(ARITHMETIC_OVERFLOW)?
                        .try_into()
                        .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?;
                    exp >>= 1;
                }

                Ok(result)
            }
        }
    }
}

pub(crate) fn book_liquidity(
    quote_type: QuoteType,
    price: u128,
    spacing: BookSpacingType,
    bid_liquidity_per_m: &[u32],
    ask_liquidity_per_m: &[u32],
    reserves_a: u64,
    reserves_b: u64,
) -> Result<SingleSideLiquidity, CoreError> {
    let liquidity_per_m = if quote_type.a_to_b() {
        bid_liquidity_per_m
    } else {
        ask_liquidity_per_m
    };

    let total_liquidity = if quote_type.a_to_b() {
        reserves_b
    } else {
        reserves_a
    };

    let total_liquidity_per_m: u64 = liquidity_per_m.iter().map(|&x| x as u64).sum();
    if total_liquidity_per_m > PER_M_DENOMINATOR as u64 {
        return Err(INVALID_ORACLE_DATA);
    }

    let mut book = SingleSideLiquidity::new();

    for (i, liquidity_share) in liquidity_per_m.iter().enumerate() {
        if *liquidity_share == 0 {
            continue;
        }

        let price = spacing.price(price, i, quote_type.a_to_b())?;

        let liquidity: u64 = u128::from(total_liquidity)
            .checked_mul(*liquidity_share as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .checked_div(PER_M_DENOMINATOR as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U64)?;

        book.push((price, liquidity));
    }

    Ok(book)
}

pub(crate) fn new_book_liquidity(
    quote: &Quote,
    price: u128,
    spacing: BookSpacingType,
    bid_liquidity_per_m: &[u32],
    ask_liquidity_per_m: &[u32],
    reserves_a: u64,
    reserves_b: u64,
) -> Result<OracleData, CoreError> {
    let mut bid_liquidity = [0u128; BOOK_LIQUIDITY_LEVELS];
    let mut ask_liquidity = [0u128; BOOK_LIQUIDITY_LEVELS];

    for (i, liquidity_per_m) in bid_liquidity_per_m.iter().enumerate() {
        bid_liquidity[i] = u128::from(*liquidity_per_m)
            .checked_mul(reserves_b as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .checked_div(PER_M_DENOMINATOR as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?;
    }

    for (i, liquidity_per_m) in ask_liquidity_per_m.iter().enumerate() {
        ask_liquidity[i] = u128::from(*liquidity_per_m)
            .checked_mul(reserves_a as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .checked_div(PER_M_DENOMINATOR as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?;
    }

    let new_reserves_a = if quote.quote_type.output_is_token_a() {
        reserves_a
            .checked_sub(quote.amount_out)
            .ok_or(ARITHMETIC_OVERFLOW)?
    } else {
        reserves_a
            .checked_add(quote.amount_in)
            .ok_or(ARITHMETIC_OVERFLOW)?
    };

    let new_reserves_b = if quote.quote_type.output_is_token_b() {
        reserves_b
            .checked_sub(quote.amount_out)
            .ok_or(ARITHMETIC_OVERFLOW)?
    } else {
        reserves_b
            .checked_add(quote.amount_in)
            .ok_or(ARITHMETIC_OVERFLOW)?
    };

    let mut remaining_consumed = quote.amount_out as u128;

    // Remove the liquidity from one side, starting from the best price
    for i in 0..BOOK_LIQUIDITY_LEVELS {
        let current_level = if quote.quote_type.a_to_b() {
            &mut bid_liquidity[i]
        } else {
            &mut ask_liquidity[i]
        };
        let current_liquidity = *current_level;
        *current_level = current_liquidity.saturating_sub(remaining_consumed);
        remaining_consumed = remaining_consumed.saturating_sub(current_liquidity);
        if remaining_consumed == 0 {
            break;
        }
    }

    // Add liquidity back to the other side at the best price
    let best_other_level = if quote.quote_type.a_to_b() {
        let best_index = ask_liquidity
            .iter_mut()
            .position(|x| *x > 0)
            .unwrap_or(BOOK_LIQUIDITY_LEVELS - 1);
        &mut ask_liquidity[best_index]
    } else {
        let best_index = bid_liquidity
            .iter_mut()
            .position(|x| *x > 0)
            .unwrap_or(BOOK_LIQUIDITY_LEVELS - 1);
        &mut bid_liquidity[best_index]
    };
    *best_other_level = best_other_level
        .checked_add(quote.amount_in as u128)
        .ok_or(ARITHMETIC_OVERFLOW)?;

    // Since the we don't know the input slice size, we explicitly allocate BOOK_LIQUIDITY_LEVELS
    // for each side
    let mut new_bid_liquidity_per_m = [0u32; BOOK_LIQUIDITY_LEVELS];
    let mut new_ask_liquidity_per_m = [0u32; BOOK_LIQUIDITY_LEVELS];

    // Calculate the new per_m liquidity for each side
    for i in 0..BOOK_LIQUIDITY_LEVELS {
        if new_reserves_b > 0 {
            new_bid_liquidity_per_m[i] = bid_liquidity[i]
                .checked_mul(PER_M_DENOMINATOR as u128)
                .ok_or(ARITHMETIC_OVERFLOW)?
                .checked_div(new_reserves_b as u128)
                .ok_or(ARITHMETIC_OVERFLOW)?
                .try_into()
                .map_err(|_| AMOUNT_EXCEEDS_MAX_U32)?;
        }
        if new_reserves_a > 0 {
            new_ask_liquidity_per_m[i] = ask_liquidity[i]
                .checked_mul(PER_M_DENOMINATOR as u128)
                .ok_or(ARITHMETIC_OVERFLOW)?
                .checked_div(new_reserves_a as u128)
                .ok_or(ARITHMETIC_OVERFLOW)?
                .try_into()
                .map_err(|_| AMOUNT_EXCEEDS_MAX_U32)?;
        }
    }

    // Normalize the liquidity to equal PER_M_DENOMINATOR by adding
    // the missing liquidity to the last price level
    let total_bid_per_m = new_bid_liquidity_per_m.iter().sum::<u32>();
    let total_ask_per_m = new_ask_liquidity_per_m.iter().sum::<u32>();

    let missing_bid_per_m = PER_M_DENOMINATOR as u32 - total_bid_per_m;
    let missing_ask_per_m = PER_M_DENOMINATOR as u32 - total_ask_per_m;

    let worst_bid_index = new_bid_liquidity_per_m
        .iter()
        .enumerate()
        .rev()
        .find(|(_, x)| **x > 0)
        .map(|(i, _)| i)
        .unwrap_or(BOOK_LIQUIDITY_LEVELS - 1);
    let worst_ask_index = new_ask_liquidity_per_m
        .iter()
        .enumerate()
        .rev()
        .find(|(_, x)| **x > 0)
        .map(|(i, _)| i)
        .unwrap_or(BOOK_LIQUIDITY_LEVELS - 1);

    new_bid_liquidity_per_m[worst_bid_index] = new_bid_liquidity_per_m[worst_bid_index]
        .checked_add(missing_bid_per_m)
        .ok_or(ARITHMETIC_OVERFLOW)?;
    new_ask_liquidity_per_m[worst_ask_index] = new_ask_liquidity_per_m[worst_ask_index]
        .checked_add(missing_ask_per_m)
        .ok_or(ARITHMETIC_OVERFLOW)?;

    Ok(OracleData::OrderBook {
        price_q64_64: price,
        spacing,
        bid_liquidity_per_m: new_bid_liquidity_per_m,
        ask_liquidity_per_m: new_ask_liquidity_per_m,
    })
}

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

    #[rstest]
    #[case(11111, true, Ok(vec![1000, 989, 978, 967, 956]))]
    #[case(11111, false, Ok(vec![1000, 1012, 1023, 1034, 1045]))]
    #[case(22222, true, Ok(vec![1000, 978, 956, 934, 912]))]
    #[case(22222, false, Ok(vec![1000, 1023, 1045, 1067, 1089]))]
    #[case(0, true, Ok(vec![1000, 1000, 1000, 1000, 1000]))]
    #[case(0, false, Ok(vec![1000, 1000, 1000, 1000, 1000]))]
    fn test_linear_spacing(
        #[case] spacing_per_m: u16,
        #[case] a_to_b: bool,
        #[case] expected: Result<Vec<u128>, CoreError>,
    ) {
        let spacing = BookSpacingType::Linear(spacing_per_m);
        let prices = (0..5)
            .map(|i| spacing.price(1000, i, a_to_b))
            .collect::<Result<Vec<u128>, CoreError>>();
        assert_eq!(prices, expected);
    }

    #[rstest]
    #[case(10000, true, Ok(vec![1000, 989, 980, 969, 960]))]
    #[case(10000, false, Ok(vec![1000, 1009, 1020, 1029, 1040]))]
    #[case(20000, true, Ok(vec![1000, 979, 960, 940, 922]))]
    #[case(20000, false, Ok(vec![1000, 1019, 1040, 1060, 1082]))]
    #[case(0, true, Ok(vec![1000, 1000, 1000, 1000, 1000]))]
    #[case(0, false, Ok(vec![1000, 1000, 1000, 1000, 1000]))]
    fn test_exponential_spacing(
        #[case] spacing_per_m: u16,
        #[case] a_to_b: bool,
        #[case] expected: Result<Vec<u128>, CoreError>,
    ) {
        let spacing = BookSpacingType::Exponential(spacing_per_m);
        let prices = (0..5)
            .map(|i| spacing.price(1000, i, a_to_b))
            .collect::<Result<Vec<u128>, CoreError>>();
        assert_eq!(prices, expected);
    }

    #[rstest]
    #[case(QuoteType::TokenAExactIn, Ok(vec![400, 1000, 600]))]
    #[case(QuoteType::TokenBExactIn, Ok(vec![300, 500, 200]))]
    #[case(QuoteType::TokenAExactOut, Ok(vec![300, 500, 200]))]
    #[case(QuoteType::TokenBExactOut, Ok(vec![400, 1000, 600]))]
    fn test_book_liquidity(
        #[case] quote_type: QuoteType,
        #[case] expected: Result<Vec<u64>, CoreError>,
    ) {
        let liquidity = book_liquidity(
            quote_type,
            1000,
            BookSpacingType::Linear(0),
            &[200_000, 0, 500_000, 300_000],
            &[300_000, 0, 500_000, 200_000],
            1000,
            2000,
        )
        .map(|liquidity| liquidity.as_slice().iter().map(|x| x.1).collect());
        assert_eq!(liquidity, expected);
    }

    #[rstest]
    #[case(QuoteType::TokenAExactIn, 100, 50, vec![0, 157894, 0, 526315, 315791], vec![0, 363636, 0, 454545, 181819])]
    #[case(QuoteType::TokenAExactIn, 50, 100, vec![0, 111111, 0, 555555, 333334], vec![0, 333333, 0, 476190, 190477])]
    #[case(QuoteType::TokenBExactIn, 100, 50, vec![0, 272727, 0, 454545, 272728], vec![0, 263157, 0, 526315, 210528])]
    #[case(QuoteType::TokenBExactIn, 50, 100, vec![0, 238095, 0, 476190, 285715], vec![0, 222222, 0, 555555, 222223])]
    #[case(QuoteType::TokenAExactIn, 500, 400, vec![0, 0, 0, 500000, 500000], vec![0, 533333, 0, 333333, 133334])]
    #[case(QuoteType::TokenAExactIn, 400, 500, vec![0, 0, 0, 400000, 600000], vec![0, 500000, 0, 357142, 142858])]
    #[case(QuoteType::TokenBExactIn, 500, 400, vec![0, 466666, 0, 333333, 200001], vec![0, 0, 0, 666666, 333334])]
    #[case(QuoteType::TokenBExactIn, 400, 500, vec![0, 428571, 0, 357142, 214287], vec![0, 0, 0, 600000, 400000])]
    #[case(QuoteType::TokenAExactIn, 0, 0, vec![0, 200_000, 0, 500_000, 300_000], vec![0, 300_000, 0, 500_000, 200_000])]
    #[case(QuoteType::TokenBExactIn, 0, 0, vec![0, 200_000, 0, 500_000, 300_000], vec![0, 300_000, 0, 500_000, 200_000])]
    fn test_new_book_liquidity(
        #[case] quote_type: QuoteType,
        #[case] amount_in: u64,
        #[case] amount_out: u64,
        #[case] expected_bid_liquidity_per_m: Vec<u32>,
        #[case] expected_ask_liquidity_per_m: Vec<u32>,
    ) {
        let price = 1 << 64;
        let spacing = BookSpacingType::Linear(0);
        let liquidity = new_book_liquidity(
            &Quote {
                amount_in,
                amount_out,
                quote_type,
            },
            price,
            spacing,
            &[0, 200_000, 0, 500_000, 300_000],
            &[0, 300_000, 0, 500_000, 200_000],
            1000,
            1000,
        )
        .unwrap();

        let mut bid_liquidity_per_m = [0u32; 32];
        let mut ask_liquidity_per_m = [0u32; 32];

        bid_liquidity_per_m[..expected_bid_liquidity_per_m.len()]
            .copy_from_slice(&expected_bid_liquidity_per_m);
        ask_liquidity_per_m[..expected_ask_liquidity_per_m.len()]
            .copy_from_slice(&expected_ask_liquidity_per_m);

        let expected = OracleData::OrderBook {
            price_q64_64: 1 << 64,
            spacing: BookSpacingType::Linear(0),
            bid_liquidity_per_m,
            ask_liquidity_per_m,
        };
        assert_eq!(liquidity, expected);
    }

    #[rstest]
    #[case(100, 50, vec![0, 157894, 0, 526315, 315791])]
    #[case(50, 100, vec![0, 111111, 0, 555555, 333334])]
    fn test_new_book_liquidity_no_a_liquidity(
        #[case] amount_in: u64,
        #[case] amount_out: u64,
        #[case] expected_bid_liquidity_per_m: Vec<u32>,
    ) {
        let price = 1 << 64;
        let spacing = BookSpacingType::Linear(0);
        let liquidity = new_book_liquidity(
            &Quote {
                amount_in,
                amount_out,
                quote_type: QuoteType::TokenAExactIn,
            },
            price,
            spacing,
            &[0, 200_000, 0, 500_000, 300_000],
            &[0, 0, 0, 0, 0],
            0,
            1000,
        )
        .unwrap();

        let mut bid_liquidity_per_m = [0u32; 32];
        let mut ask_liquidity_per_m = [0u32; 32];

        bid_liquidity_per_m[..expected_bid_liquidity_per_m.len()]
            .copy_from_slice(&expected_bid_liquidity_per_m);
        ask_liquidity_per_m[BOOK_LIQUIDITY_LEVELS - 1] = PER_M_DENOMINATOR as u32;

        let expected = OracleData::OrderBook {
            price_q64_64: 1 << 64,
            spacing: BookSpacingType::Linear(0),
            bid_liquidity_per_m,
            ask_liquidity_per_m,
        };
        assert_eq!(liquidity, expected);
    }

    #[rstest]
    #[case(100, 50, vec![0, 263157, 0, 526315, 210528])]
    #[case(50, 100, vec![0, 222222, 0, 555555, 222223])]
    fn test_new_book_liquidity_no_b_liquidity(
        #[case] amount_in: u64,
        #[case] amount_out: u64,
        #[case] expected_ask_liquidity_per_m: Vec<u32>,
    ) {
        let price = 1 << 64;
        let spacing = BookSpacingType::Linear(0);
        let liquidity = new_book_liquidity(
            &Quote {
                amount_in,
                amount_out,
                quote_type: QuoteType::TokenBExactIn,
            },
            price,
            spacing,
            &[0, 0, 0, 0, 0],
            &[0, 300_000, 0, 500_000, 200_000],
            1000,
            0,
        )
        .unwrap();

        let mut bid_liquidity_per_m = [0u32; 32];
        let mut ask_liquidity_per_m = [0u32; 32];

        bid_liquidity_per_m[BOOK_LIQUIDITY_LEVELS - 1] = PER_M_DENOMINATOR as u32;
        ask_liquidity_per_m[..expected_ask_liquidity_per_m.len()]
            .copy_from_slice(&expected_ask_liquidity_per_m);

        let expected = OracleData::OrderBook {
            price_q64_64: 1 << 64,
            spacing: BookSpacingType::Linear(0),
            bid_liquidity_per_m,
            ask_liquidity_per_m,
        };
        assert_eq!(liquidity, expected);
    }
}