riptide-amm-math 1.0.4

use ethnum::U256;
#[cfg(feature = "floats")]
use libm::pow;

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

use super::{AMOUNT_EXCEEDS_MAX_I32, PER_M_DENOMINATOR};

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

use super::U128;

#[cfg(feature = "floats")]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn amount_to_ui_amount(amount: u64, decimals: u8) -> f64 {
    let power = pow(10f64, decimals as f64);
    amount as f64 / power
}

#[cfg(feature = "floats")]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn ui_amount_to_amount(amount: f64, decimals: u8) -> u64 {
    let power = pow(10f64, decimals as f64);
    (amount * power) as u64
}

/// Convert an amount in token A to an amount in token B
///
/// # Parameters
/// * `amount_a` - The amount in token A
/// * `price` - The Q64.64 price in B/A
/// * `round_up` - Whether to round up the result
///
/// # Returns
/// * `u64` - The amount in token B
#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn a_to_b(amount_a: u64, price: U128, round_up: bool) -> Result<u64, CoreError> {
    #[allow(clippy::useless_conversion)]
    let price: u128 = price.into();

    let product = u128::from(amount_a)
        .checked_mul(price)
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let quotient = product >> 64;
    let remainder = product as u64;

    let result = if round_up && remainder > 0 {
        quotient + 1
    } else {
        quotient
    };

    result.try_into().map_err(|_| AMOUNT_EXCEEDS_MAX_U64)
}

/// Convert an amount in token B to an amount in token A
///
/// # Parameters
/// * `amount_b` - The amount in token B
/// * `price` - The Q64.64 price in B/A
/// * `round_up` - Whether to round up the result
///
/// # Returns
/// * `u64` - The amount in token A
#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn b_to_a(amount_b: u64, price: U128, round_up: bool) -> Result<u64, CoreError> {
    #[allow(clippy::useless_conversion)]
    let price: u128 = price.into();
    if price == 0 {
        return Ok(0);
    }

    let numerator = u128::from(amount_b)
        .checked_shl(64)
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let quotient = numerator.checked_div(price).ok_or(ARITHMETIC_OVERFLOW)?;
    let remainder = numerator.checked_rem(price).ok_or(ARITHMETIC_OVERFLOW)?;

    let result = if round_up && remainder > 0 {
        quotient + 1
    } else {
        quotient
    };

    result.try_into().map_err(|_| AMOUNT_EXCEEDS_MAX_U64)
}

/// Computes deviation_per_m from inventory ratio.
/// 0 = balanced (50/50), +1_000_000 = 100% token_a, -1_000_000 = 100% token_b.
/// inventory_ratio = value_a / (value_a + value_b) where values are in q64.64.
#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn deviation_per_m(
    price_q64_64: U128,
    reserves_a: u64,
    reserves_b: u64,
) -> Result<i32, CoreError> {
    let value_a = U256::from(reserves_a as u128)
        .checked_mul(U256::from(price_q64_64))
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let value_b = U256::from(reserves_b as u128)
        .checked_shl(64)
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let total = value_a.checked_add(value_b).ok_or(ARITHMETIC_OVERFLOW)?;
    if total == U256::ZERO {
        return Ok(0);
    }
    let twice_value_a = value_a
        .checked_mul(U256::from(2u128))
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let ratio_per_m = twice_value_a
        .checked_mul(U256::from(PER_M_DENOMINATOR as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?
        .checked_div(total)
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let ratio: i32 = ratio_per_m.try_into().map_err(|_| AMOUNT_EXCEEDS_MAX_I32)?;
    ratio
        .checked_sub(PER_M_DENOMINATOR)
        .ok_or(ARITHMETIC_OVERFLOW)
}

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

    #[cfg(feature = "floats")]
    #[rstest]
    #[case(1000000000, 9, 1.0)]
    #[case(1000000000, 6, 1000.0)]
    #[case(1000000000, 3, 1000000.0)]
    fn test_amount_to_ui_amount(
        #[case] amount: u64,
        #[case] decimals: u8,
        #[case] expected_ui_amount: f64,
    ) {
        let ui_amount = amount_to_ui_amount(amount, decimals);
        assert_eq!(ui_amount, expected_ui_amount);
    }

    #[cfg(feature = "floats")]
    #[rstest]
    #[case(1.0, 9, 1000000000)]
    #[case(1000.0, 6, 1000000000)]
    #[case(1000000.0, 3, 1000000000)]
    fn test_ui_amount_to_amount(
        #[case] ui_amount: f64,
        #[case] decimals: u8,
        #[case] expected_amount: u64,
    ) {
        let amount = ui_amount_to_amount(ui_amount, decimals);
        assert_eq!(amount, expected_amount);
    }

    #[rstest]
    #[case(100, 1 << 64, true, Ok(100))]
    #[case(100, 1 << 64, false, Ok(100))]
    #[case(100, 8 << 64, true, Ok(800))]
    #[case(100, 8 << 64, false, Ok(800))]
    #[case(100, (1 << 64) / 8, true, Ok(13))]
    #[case(100, (1 << 64) / 8, false, Ok(12))]
    #[case(100, 0, true, Ok(0))]
    #[case(0, 1 << 64, true, Ok(0))]
    fn test_a_to_b(
        #[case] amount_a: u64,
        #[case] price: u128,
        #[case] round_up: bool,
        #[case] expected: Result<u64, CoreError>,
    ) {
        let result = a_to_b(amount_a, U128::from(price), round_up);
        assert_eq!(result, expected);
    }

    #[rstest]
    #[case(100, 1 << 64, true, Ok(100))]
    #[case(100, 1 << 64, false, Ok(100))]
    #[case(100, 8 << 64, true, Ok(13))]
    #[case(100, 8 << 64, false, Ok(12))]
    #[case(100, (1 << 64) / 8, true, Ok(800))]
    #[case(100, (1 << 64) / 8, false, Ok(800))]
    #[case(100, 0, true, Ok(0))]
    #[case(0, 1 << 64, true, Ok(0))]
    fn test_b_to_a(
        #[case] amount_b: u64,
        #[case] price: u128,
        #[case] round_up: bool,
        #[case] expected: Result<u64, CoreError>,
    ) {
        let result = b_to_a(amount_b, U128::from(price), round_up);
        assert_eq!(result, expected);
    }

    #[rstest]
    #[case(1 << 64, 500, 500, 0)] // 50/50 -> 0
    #[case(1 << 64, 750, 250, 500_000)] // 75% token_a -> +500_000
    #[case(1 << 64, 250, 750, -500_000)] // 25% token_a -> -500_000
    #[case(1 << 64, 1000, 0, 1_000_000)] // 100% token_a -> +1_000_000
    #[case(1 << 64, 0, 1000, -1_000_000)] // 0% token_a -> -1_000_000
    fn test_deviation_per_m(
        #[case] price: u128,
        #[case] reserves_a: u64,
        #[case] reserves_b: u64,
        #[case] expected: i32,
    ) {
        let result = deviation_per_m(U128::from(price), reserves_a, reserves_b).unwrap();
        assert_eq!(result, expected);
    }

    #[test]
    fn test_deviation_zero_reserves() {
        assert_eq!(deviation_per_m(U128::from(1u128 << 64), 0, 0).unwrap(), 0);
    }
}