equanetwork-math 0.1.0

use ethnum::U256;

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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OraclePrice {
    pub price: u128,
    // Confidence in price as a u128 confidence interval (lower is more confident)
    pub confidence: u128,
    pub timestamp: i64,
}

impl OraclePrice {
    pub fn try_from_raw(
        price: i64,
        confidence: u64,
        exponent: i32,
        timestamp: i64,
    ) -> Result<OraclePrice, CoreError> {
        Self::try_from_scaled(
            i128::from(price),
            u128::from(confidence),
            exponent,
            timestamp,
        )
    }

    /// Convert a signed mantissa (Pyth i64, Switchboard/Chainlink i128) into Q64.64.
    pub fn try_from_scaled(
        price: i128,
        confidence: u128,
        exponent: i32,
        timestamp: i64,
    ) -> Result<OraclePrice, CoreError> {
        if price <= 0 {
            return Err(ARITHMETIC_OVERFLOW);
        }
        let price = mantissa_to_q64(price as u128, exponent)?;
        if price == 0 {
            return Err(ARITHMETIC_OVERFLOW);
        }
        let confidence = mantissa_to_q64(confidence, exponent)?;
        Ok(OraclePrice {
            price,
            confidence,
            timestamp,
        })
    }

    /// Pair rate `in / out` in atomic B/A Q64.64, using the confidence bounds
    /// that minimize output tokens per input token, then scaling by
    /// `10^(decimals_out - decimals_in)` with floor division.
    pub fn combine(
        price_in: &OraclePrice,
        price_out: &OraclePrice,
        decimals_in: u8,
        decimals_out: u8,
    ) -> Result<u128, CoreError> {
        let in_lo = price_in.price.saturating_sub(price_in.confidence);
        let out_hi = price_out.price.saturating_add(price_out.confidence);
        if in_lo == 0 || out_hi == 0 {
            return Err(ARITHMETIC_OVERFLOW);
        }

        let ui_rate: u128 = U256::from(in_lo)
            .checked_shl(64)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .checked_div(U256::from(out_hi))
            .ok_or(ARITHMETIC_OVERFLOW)?
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?;
        if ui_rate == 0 {
            return Err(ARITHMETIC_OVERFLOW);
        }

        let scaled = if decimals_out >= decimals_in {
            let scale = 10u128
                .checked_pow((decimals_out - decimals_in) as u32)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            ui_rate.checked_mul(scale).ok_or(ARITHMETIC_OVERFLOW)?
        } else {
            let scale = 10u128
                .checked_pow((decimals_in - decimals_out) as u32)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            ui_rate / scale
        };
        if scaled == 0 {
            return Err(ARITHMETIC_OVERFLOW);
        }
        Ok(scaled)
    }
}

fn mantissa_to_q64(mantissa: u128, exponent: i32) -> Result<u128, CoreError> {
    let scale = 10u128
        .checked_pow(exponent.unsigned_abs())
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let shifted = U256::from(mantissa)
        .checked_shl(64)
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let result = if exponent >= 0 {
        shifted
            .checked_mul(U256::from(scale))
            .ok_or(ARITHMETIC_OVERFLOW)?
    } else {
        shifted
            .checked_div(U256::from(scale))
            .ok_or(ARITHMETIC_OVERFLOW)?
    };
    result.try_into().map_err(|_| AMOUNT_EXCEEDS_MAX_U128)
}

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

    fn price(mid: u128, conf: u128) -> OraclePrice {
        OraclePrice {
            price: mid,
            confidence: conf,
            timestamp: 0,
        }
    }

    #[rstest]
    #[case(100_000_000, 0, -8, 1_700_000_000, Ok((1u128 << 64, 0)))]
    #[case(200_000_000, 1_000_000, -8, 42, Ok((2 << 64, (1u128 << 64) / 100)))]
    #[case(1, 0, 0, 0, Ok((1u128 << 64, 0)))]
    #[case(1, 0, 1, 0, Ok((10 << 64, 0)))]
    #[case(1, 1, -1, 7, Ok(((1u128 << 64) / 10, (1u128 << 64) / 10)))]
    #[case(0, 0, -8, 0, Err(ARITHMETIC_OVERFLOW))]
    #[case(-1, 0, -8, 0, Err(ARITHMETIC_OVERFLOW))]
    #[case(1, 0, -20, 0, Err(ARITHMETIC_OVERFLOW))]
    fn test_try_from_raw(
        #[case] raw_price: i64,
        #[case] raw_confidence: u64,
        #[case] exponent: i32,
        #[case] timestamp: i64,
        #[case] expected: Result<(u128, u128), CoreError>,
    ) {
        let result = OraclePrice::try_from_raw(raw_price, raw_confidence, exponent, timestamp);
        match (result, expected) {
            (Ok(obs), Ok((price, confidence))) => {
                let got_price: u128 = obs.price.into();
                let got_confidence: u128 = obs.confidence.into();
                assert_eq!(got_price, price);
                assert_eq!(got_confidence, confidence);
                assert_eq!(obs.timestamp, timestamp);
            }
            (Err(err), Err(expected_err)) => assert_eq!(err, expected_err),
            (result, expected) => panic!("unexpected result {result:?}, expected {expected:?}"),
        }
    }

    #[test]
    fn try_from_scaled_switchboard_precision() {
        // 1.0 with Switchboard On-Demand PRECISION=18 must fit via U256.
        let obs = OraclePrice::try_from_scaled(1_000_000_000_000_000_000, 0, -18, 7).unwrap();
        assert_eq!(obs.price, 1u128 << 64);
        assert_eq!(obs.confidence, 0);
        assert_eq!(obs.timestamp, 7);
    }

    #[rstest]
    #[case(1u128 << 64, 0, 1u128 << 64, 0, 6, 6, Ok(1u128 << 64))]
    #[case(2 << 64, 0, 1u128 << 64, 0, 6, 6, Ok(2 << 64))]
    #[case(1u128 << 64, 0, 2 << 64, 0, 9, 9, Ok((1u128 << 64) / 2))]
    #[case(200 << 64, 0, 100 << 64, 0, 0, 0, Ok(2 << 64))]
    #[case(1u128 << 64, (1u128 << 64) / 16, 1u128 << 64, (1u128 << 64) / 16, 6, 6, Ok((15 << 64) / 17))]
    #[case(1u128 << 64, 1u128 << 64, 1u128 << 64, 0, 6, 6, Err(ARITHMETIC_OVERFLOW))]
    #[case(0, 0, 1u128 << 64, 0, 6, 6, Err(ARITHMETIC_OVERFLOW))]
    #[case(1u128 << 64, 0, 0, 0, 6, 6, Err(ARITHMETIC_OVERFLOW))]
    #[case(1u128 << 64, 0, 1u128 << 64, 0, 6, 9, Ok((1u128 << 64) * 1000))]
    #[case(1u128 << 64, 0, 1u128 << 64, 0, 9, 6, Ok((1u128 << 64) / 1000))]
    fn test_combine(
        #[case] in_price: u128,
        #[case] in_conf: u128,
        #[case] out_price: u128,
        #[case] out_conf: u128,
        #[case] decimals_in: u8,
        #[case] decimals_out: u8,
        #[case] expected: Result<u128, CoreError>,
    ) {
        let result = OraclePrice::combine(
            &price(in_price, in_conf),
            &price(out_price, out_conf),
            decimals_in,
            decimals_out,
        );
        assert_eq!(result, expected);
    }

    #[test]
    fn combine_floor_not_ceil_when_downscaling() {
        let one = 1u128 << 64;
        let floored = OraclePrice::combine(&price(one, 0), &price(one, 0), 9, 6).unwrap();
        assert_eq!(floored, one / 1000);
        assert_ne!(floored, one.div_ceil(1000));
    }

    #[test]
    fn combine_uses_worst_bound_for_the_user() {
        let mid =
            OraclePrice::combine(&price(1u128 << 64, 0), &price(1u128 << 64, 0), 6, 6).unwrap();
        let conf = (1u128 << 64) / 16;
        let conservative =
            OraclePrice::combine(&price(1u128 << 64, conf), &price(1u128 << 64, conf), 6, 6)
                .unwrap();
        let optimistic_in =
            OraclePrice::combine(&price(1u128 << 64, 0), &price(1u128 << 64, conf), 6, 6).unwrap();
        let optimistic_out =
            OraclePrice::combine(&price(1u128 << 64, conf), &price(1u128 << 64, 0), 6, 6).unwrap();

        assert!(conservative < mid);
        assert_eq!(conservative, (15 << 64) / 17);
        assert_eq!(optimistic_in, (1u128 << 64) * 16 / 17);
        assert_eq!(optimistic_out, (1u128 << 64) * 15 / 16);
        assert!(conservative < optimistic_in);
        assert!(conservative < optimistic_out);
    }
}