equanetwork-math 0.1.0

use super::{
    a_to_b, b_to_a, fee_from_post_fee_amount, fee_from_pre_fee_amount, next_band, prev_band,
    skew_band, CoreError, SkewBand, SkewExponent, SkewVault, ARITHMETIC_OVERFLOW,
    PARTIAL_FILL_NOT_ALLOWED, PER_M_DENOMINATOR, U128,
};

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

use ethnum::U256;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub enum SwapMode {
    ExactIn,
    ExactOut,
}

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

#[derive(Debug, Clone, Copy)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub struct SwapVault {
    pub swap_fee_per_m: u32,
    pub inventory_limit_upper: u64,
    pub inventory_limit_lower: u64,
    pub max_swap_amount: u64,
    pub positive_skew_per_m: u32,
    pub negative_skew_per_m: u32,
    pub lower_skew_offset: u64,
    pub upper_skew_offset: u64,
    pub skew_exponent: SkewExponent,
}

#[allow(clippy::too_many_arguments)]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub fn swap_quote(
    amount: u64,
    swap_mode: SwapMode,
    base_price: U128,
    allow_partial_fill: bool,
    input_vault: SwapVault,
    input_vault_balance: u64,
    output_vault: SwapVault,
    output_vault_balance: u64,
) -> Result<SwapQuote, CoreError> {
    let base_price: u128 = base_price.into();
    let max_amount_in = input_vault
        .inventory_limit_upper
        .checked_sub(input_vault_balance)
        .ok_or(ARITHMETIC_OVERFLOW)?
        .min(input_vault.max_swap_amount);

    let max_amount_out = output_vault_balance
        .checked_sub(output_vault.inventory_limit_lower)
        .ok_or(ARITHMETIC_OVERFLOW)?
        .min(output_vault.max_swap_amount);

    match swap_mode {
        SwapMode::ExactIn => {
            let amount_in = amount.min(max_amount_in);
            let mut quote = raw_quote(
                amount_in,
                SwapMode::ExactIn,
                base_price.into(),
                &input_vault,
                input_vault_balance,
                &output_vault,
                output_vault_balance,
            )?;
            if quote.amount_out > max_amount_out {
                quote = raw_quote(
                    max_amount_out,
                    SwapMode::ExactOut,
                    base_price.into(),
                    &input_vault,
                    input_vault_balance,
                    &output_vault,
                    output_vault_balance,
                )?;
            }
            if allow_partial_fill || quote.amount_in == amount {
                Ok(quote)
            } else {
                Err(PARTIAL_FILL_NOT_ALLOWED)
            }
        }
        SwapMode::ExactOut => {
            let amount_out = amount.min(max_amount_out);
            let mut quote = raw_quote(
                amount_out,
                SwapMode::ExactOut,
                base_price.into(),
                &input_vault,
                input_vault_balance,
                &output_vault,
                output_vault_balance,
            )?;
            if quote.amount_in > max_amount_in {
                quote = raw_quote(
                    max_amount_in,
                    SwapMode::ExactIn,
                    base_price.into(),
                    &input_vault,
                    input_vault_balance,
                    &output_vault,
                    output_vault_balance,
                )?;
            }
            if allow_partial_fill || quote.amount_out == amount {
                Ok(quote)
            } else {
                Err(PARTIAL_FILL_NOT_ALLOWED)
            }
        }
    }
}

fn raw_quote(
    amount: u64,
    swap_mode: SwapMode,
    base_price: u128,
    input_vault: &SwapVault,
    input_vault_balance: u64,
    output_vault: &SwapVault,
    output_vault_balance: u64,
) -> Result<SwapQuote, CoreError> {
    let prelim_fee = match swap_mode {
        SwapMode::ExactIn => fee_from_pre_fee_amount(amount, input_vault.swap_fee_per_m)?,
        SwapMode::ExactOut => fee_from_post_fee_amount(amount, output_vault.swap_fee_per_m)?,
    };

    let consumable_amount = match swap_mode {
        SwapMode::ExactIn => amount.checked_sub(prelim_fee).ok_or(ARITHMETIC_OVERFLOW)?,
        SwapMode::ExactOut => amount.checked_add(prelim_fee).ok_or(ARITHMETIC_OVERFLOW)?,
    };

    let mut raw_amount_in = 0u64;
    let mut raw_amount_out = 0u64;
    let mut remaining_amount = consumable_amount;

    let mut current_input_balance = input_vault_balance;
    let mut current_output_balance = output_vault_balance;

    let mut input_band = Some(skew_band(
        input_vault_balance,
        true,
        SkewVault::from(*input_vault),
    )?);
    let mut output_band = Some(skew_band(
        output_vault_balance,
        false,
        SkewVault::from(*output_vault),
    )?);

    while remaining_amount > 0 {
        let (Some(current_input_band), Some(current_output_band)) = (&input_band, &output_band)
        else {
            break;
        };

        let (a_to_b_price, b_to_a_price) =
            band_prices(base_price, current_input_band, current_output_band)?;
        if a_to_b_price == 0 || b_to_a_price == 0 {
            break;
        }

        let (input_amount, output_amount, move_input_vault, move_output_vault) = match swap_mode {
            SwapMode::ExactIn => {
                let input_amount_input_band = current_input_band
                    .upper_limit
                    .abs_diff(current_input_balance)
                    .min(remaining_amount);
                let output_amount_output_band = current_output_band
                    .lower_limit
                    .abs_diff(current_output_balance);
                let input_amount_output_band =
                    b_to_a(output_amount_output_band, b_to_a_price.into(), true)?
                        .min(remaining_amount);

                let input_amount = input_amount_input_band.min(input_amount_output_band);
                let output_amount = a_to_b(input_amount, a_to_b_price.into(), false)?;

                (
                    input_amount,
                    output_amount,
                    input_amount_input_band < input_amount_output_band,
                    input_amount_input_band > input_amount_output_band,
                )
            }
            SwapMode::ExactOut => {
                let output_amount_output_band = current_output_band
                    .lower_limit
                    .abs_diff(current_output_balance)
                    .min(remaining_amount);
                let input_amount_input_band = current_input_band
                    .upper_limit
                    .abs_diff(current_input_balance);
                let output_amount_input_band =
                    a_to_b(input_amount_input_band, a_to_b_price.into(), false)?
                        .min(remaining_amount);

                let output_amount = output_amount_input_band.min(output_amount_output_band);
                let input_amount = b_to_a(output_amount, b_to_a_price.into(), true)?;

                (
                    input_amount,
                    output_amount,
                    output_amount_input_band < output_amount_output_band,
                    output_amount_input_band > output_amount_output_band,
                )
            }
        };

        raw_amount_in = raw_amount_in
            .checked_add(input_amount)
            .ok_or(ARITHMETIC_OVERFLOW)?;
        raw_amount_out = raw_amount_out
            .checked_add(output_amount)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        remaining_amount = match swap_mode {
            SwapMode::ExactIn => remaining_amount
                .checked_sub(input_amount)
                .ok_or(ARITHMETIC_OVERFLOW)?,
            SwapMode::ExactOut => remaining_amount
                .checked_sub(output_amount)
                .ok_or(ARITHMETIC_OVERFLOW)?,
        };

        current_input_balance = current_input_balance
            .checked_add(input_amount)
            .ok_or(ARITHMETIC_OVERFLOW)?;
        current_output_balance = current_output_balance
            .checked_sub(output_amount)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        if move_input_vault {
            input_band = next_band(current_input_band, &SkewVault::from(*input_vault))?;
        }

        if move_output_vault {
            output_band = prev_band(current_output_band, &SkewVault::from(*output_vault))?;
        }
    }

    let (fee_in, amount_in) =
        if matches!(swap_mode, SwapMode::ExactIn) && consumable_amount == raw_amount_in {
            (prelim_fee, amount)
        } else {
            let fee_in = fee_from_post_fee_amount(raw_amount_in, input_vault.swap_fee_per_m)?;
            let amount_in = raw_amount_in
                .checked_add(fee_in)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            (fee_in, amount_in)
        };

    let (fee_out, amount_out) =
        if matches!(swap_mode, SwapMode::ExactOut) && consumable_amount == raw_amount_out {
            (prelim_fee, amount)
        } else {
            let fee_out = fee_from_pre_fee_amount(raw_amount_out, output_vault.swap_fee_per_m)?;
            let amount_out = raw_amount_out
                .checked_sub(fee_out)
                .ok_or(ARITHMETIC_OVERFLOW)?;
            (fee_out, amount_out)
        };

    Ok(SwapQuote {
        amount_in,
        fee_in,
        amount_out,
        fee_out,
    })
}

fn band_prices(
    base_price: u128,
    input_band: &SkewBand,
    output_band: &SkewBand,
) -> Result<(u128, u128), CoreError> {
    let price_skew_per_m = PER_M_DENOMINATOR
        .checked_add(input_band.skew_per_m)
        .ok_or(ARITHMETIC_OVERFLOW)?
        .checked_add(output_band.skew_per_m)
        .ok_or(ARITHMETIC_OVERFLOW)?
        .clamp(0, PER_M_DENOMINATOR * 2);

    let product = U256::from(base_price)
        .checked_mul(U256::from(price_skew_per_m as u32))
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let quotient = product
        .checked_div(U256::from(PER_M_DENOMINATOR as u32))
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let remainder = product
        .checked_rem(U256::from(PER_M_DENOMINATOR as u32))
        .ok_or(ARITHMETIC_OVERFLOW)?;

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

    Ok((
        quotient.try_into().map_err(|_| ARITHMETIC_OVERFLOW)?,
        result.try_into().map_err(|_| ARITHMETIC_OVERFLOW)?,
    ))
}

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

    fn stub_band(skew_per_m: i32) -> SkewBand {
        SkewBand {
            index: 0,
            lower_limit: 0,
            upper_limit: 0,
            skew_per_m,
        }
    }

    fn dead_zone_vault(swap_fee_per_m: u32, max_swap_amount: u64) -> SwapVault {
        SwapVault {
            swap_fee_per_m,
            inventory_limit_upper: 100_000,
            inventory_limit_lower: 0,
            max_swap_amount,
            positive_skew_per_m: 0,
            negative_skew_per_m: 0,
            lower_skew_offset: 0,
            upper_skew_offset: 100_000,
            skew_exponent: SkewExponent::Linear,
        }
    }

    fn skew_vault(
        positive_skew_per_m: u32,
        negative_skew_per_m: u32,
        skew_exponent: SkewExponent,
    ) -> SwapVault {
        SwapVault {
            swap_fee_per_m: 0,
            inventory_limit_upper: 64_000,
            inventory_limit_lower: 0,
            max_swap_amount: u64::MAX,
            positive_skew_per_m,
            negative_skew_per_m,
            lower_skew_offset: 16_000,
            upper_skew_offset: 48_000,
            skew_exponent,
        }
    }

    #[rstest]
    // base = PER_M so factor math is exact (no Q64.64 remainder)
    #[case(0, 0, 1_000_000u128, 1_000_000u128)]
    #[case(100_000, 0, 1_100_000, 1_100_000)]
    #[case(0, -100_000, 900_000, 900_000)]
    #[case(-600_000, -600_000, 0, 0)]
    #[case(600_000, 600_000, 2_000_000, 2_000_000)]
    fn test_band_prices(
        #[case] input_skew_per_m: i32,
        #[case] output_skew_per_m: i32,
        #[case] expected_a_to_b: u128,
        #[case] expected_b_to_a: u128,
    ) {
        let (a_to_b_price, b_to_a_price) = band_prices(
            PER_M_DENOMINATOR as u128,
            &stub_band(input_skew_per_m),
            &stub_band(output_skew_per_m),
        )
        .unwrap();
        assert_eq!(a_to_b_price, expected_a_to_b);
        assert_eq!(b_to_a_price, expected_b_to_a);
    }

    #[rstest]
    // (1<<64)*(PER_M+1) is not divisible by PER_M → floor vs ceil split
    #[case(1, 0, 18_446_762_520_453_625_325, 18_446_762_520_453_625_326)]
    fn test_band_prices_rounding(
        #[case] input_skew_per_m: i32,
        #[case] output_skew_per_m: i32,
        #[case] expected_a_to_b: u128,
        #[case] expected_b_to_a: u128,
    ) {
        let (a_to_b_price, b_to_a_price) = band_prices(
            1u128 << 64,
            &stub_band(input_skew_per_m),
            &stub_band(output_skew_per_m),
        )
        .unwrap();
        assert_eq!(a_to_b_price, expected_a_to_b);
        assert_eq!(b_to_a_price, expected_b_to_a);
        assert!(b_to_a_price > a_to_b_price);
    }

    #[rstest]
    #[case(1_000, SwapMode::ExactIn)]
    #[case(1_000, SwapMode::ExactOut)]
    #[case(0, SwapMode::ExactIn)]
    #[case(0, SwapMode::ExactOut)]
    fn test_swap_quote_identity(#[case] amount: u64, #[case] swap_mode: SwapMode) {
        let vault = dead_zone_vault(0, u64::MAX);
        let quote = swap_quote(
            amount,
            swap_mode,
            U128::from(1u128 << 64),
            true,
            vault,
            50_000,
            vault,
            50_000,
        )
        .unwrap();
        assert_eq!(
            quote,
            SwapQuote {
                amount_in: amount,
                fee_in: 0,
                amount_out: amount,
                fee_out: 0,
            }
        );
    }

    #[rstest]
    // ExactIn 10% input fee: fee_in=100, consumable 900 → out 900
    #[case(SwapMode::ExactIn, 100_000, 0, 1_000, 1_000, 100, 900, 0)]
    // ExactOut 10% output fee: amount_out=900, fee_out=100, in=1000
    #[case(SwapMode::ExactOut, 0, 100_000, 900, 1_000, 0, 900, 100)]
    // Both sides 10%: ExactIn 1000 → fee_in=100, raw 900, fee_out=90, out=810
    #[case(SwapMode::ExactIn, 100_000, 100_000, 1_000, 1_000, 100, 810, 90)]
    fn test_swap_quote_fees(
        #[case] swap_mode: SwapMode,
        #[case] input_fee_per_m: u32,
        #[case] output_fee_per_m: u32,
        #[case] amount: u64,
        #[case] expected_in: u64,
        #[case] expected_fee_in: u64,
        #[case] expected_out: u64,
        #[case] expected_fee_out: u64,
    ) {
        let quote = swap_quote(
            amount,
            swap_mode,
            U128::from(1u128 << 64),
            true,
            dead_zone_vault(input_fee_per_m, u64::MAX),
            50_000,
            dead_zone_vault(output_fee_per_m, u64::MAX),
            50_000,
        )
        .unwrap();
        assert_eq!(quote.amount_in, expected_in);
        assert_eq!(quote.fee_in, expected_fee_in);
        assert_eq!(quote.amount_out, expected_out);
        assert_eq!(quote.fee_out, expected_fee_out);
    }

    #[rstest]
    #[case(9, 100, 9, 1, 8, 0)]
    fn test_swap_quote_fees_rounding(
        #[case] amount: u64,
        #[case] input_fee_per_m: u32,
        #[case] expected_in: u64,
        #[case] expected_fee_in: u64,
        #[case] expected_out: u64,
        #[case] expected_fee_out: u64,
    ) {
        let quote = swap_quote(
            amount,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            dead_zone_vault(input_fee_per_m, u64::MAX),
            50_000,
            dead_zone_vault(0, u64::MAX),
            50_000,
        )
        .unwrap();
        assert_eq!(quote.amount_in, expected_in);
        assert_eq!(quote.fee_in, expected_fee_in);
        assert_eq!(quote.amount_out, expected_out);
        assert_eq!(quote.fee_out, expected_fee_out);
    }

    #[rstest]
    // Input at band +16 (balance 56000), increasing away → skew +16000 → out 406 for 400 in
    #[case(55_501, 32_000, true, 400, 400, 406)]
    // Output at band -16 (balance 8000), decreasing away → skew +16000 → out 406
    #[case(50_000, 8_000, false, 400, 400, 406)]
    // Dead zone control
    #[case(32_000, 32_000, true, 400, 400, 400)]
    fn test_swap_quote_skew(
        #[case] input_balance: u64,
        #[case] output_balance: u64,
        #[case] skew_on_input: bool,
        #[case] amount: u64,
        #[case] expected_in: u64,
        #[case] expected_out: u64,
    ) {
        let (input_vault, output_vault) = if skew_on_input {
            (
                skew_vault(32_000, 0, SkewExponent::Linear),
                dead_zone_vault(0, u64::MAX),
            )
        } else if input_balance == 32_000 && output_balance == 32_000 {
            (dead_zone_vault(0, u64::MAX), dead_zone_vault(0, u64::MAX))
        } else {
            (
                dead_zone_vault(0, u64::MAX),
                skew_vault(32_000, 0, SkewExponent::Linear),
            )
        };
        let quote = swap_quote(
            amount,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            input_vault,
            input_balance,
            output_vault,
            output_balance,
        )
        .unwrap();
        assert_eq!(quote.amount_in, expected_in);
        assert_eq!(quote.fee_in, 0);
        assert_eq!(quote.amount_out, expected_out);
        assert_eq!(quote.fee_out, 0);
    }

    #[rstest]
    // Band +1 skew ceil(10000/32)=313; 10000 crosses bands → path-averaged out
    #[case(48_001, 50_000, 10_000, 10_000, 10_023)]
    fn test_swap_quote_skew_rounding(
        #[case] input_balance: u64,
        #[case] output_balance: u64,
        #[case] positive_skew_per_m: u32,
        #[case] amount: u64,
        #[case] expected_out: u64,
    ) {
        let quote = swap_quote(
            amount,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            skew_vault(positive_skew_per_m, 0, SkewExponent::Linear),
            input_balance,
            dead_zone_vault(0, u64::MAX),
            output_balance,
        )
        .unwrap();
        assert_eq!(quote.amount_in, amount);
        assert_eq!(quote.amount_out, expected_out);
    }

    #[rstest]
    // ExactIn: input upper 10_000, balance 9_000 → room 1_000
    #[case(
        SwapMode::ExactIn,
        5_000,
        9_000,
        10_000,
        50_000,
        100_000,
        u64::MAX,
        1_000,
        1_000
    )]
    // ExactOut: output lower 0, balance 1_000 → room 1_000
    #[case(
        SwapMode::ExactOut,
        5_000,
        50_000,
        100_000,
        1_000,
        100_000,
        u64::MAX,
        1_000,
        1_000
    )]
    // max_swap_amount 500 caps ExactIn
    #[case(
        SwapMode::ExactIn,
        5_000,
        50_000,
        100_000,
        50_000,
        100_000,
        500,
        500,
        500
    )]
    fn test_swap_quote_caps(
        #[case] swap_mode: SwapMode,
        #[case] amount: u64,
        #[case] input_balance: u64,
        #[case] input_upper: u64,
        #[case] output_balance: u64,
        #[case] output_upper: u64,
        #[case] max_swap_amount: u64,
        #[case] expected_in: u64,
        #[case] expected_out: u64,
    ) {
        let input = SwapVault {
            swap_fee_per_m: 0,
            inventory_limit_upper: input_upper,
            inventory_limit_lower: 0,
            max_swap_amount,
            positive_skew_per_m: 0,
            negative_skew_per_m: 0,
            lower_skew_offset: 0,
            upper_skew_offset: input_upper,
            skew_exponent: SkewExponent::Linear,
        };
        let output = SwapVault {
            swap_fee_per_m: 0,
            inventory_limit_upper: output_upper,
            inventory_limit_lower: 0,
            max_swap_amount,
            positive_skew_per_m: 0,
            negative_skew_per_m: 0,
            lower_skew_offset: 0,
            upper_skew_offset: output_upper,
            skew_exponent: SkewExponent::Linear,
        };
        let quote = swap_quote(
            amount,
            swap_mode,
            U128::from(1u128 << 64),
            true,
            input,
            input_balance,
            output,
            output_balance,
        )
        .unwrap();
        assert_eq!(quote.amount_in, expected_in);
        assert_eq!(quote.amount_out, expected_out);
    }

    #[test]
    fn test_swap_quote_caps_requote_on_output_room() {
        // Input room large; output only has 500 to lower bound → ExactIn requotes via ExactOut
        let quote = swap_quote(
            5_000,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            dead_zone_vault(0, u64::MAX),
            50_000,
            dead_zone_vault(0, u64::MAX),
            500,
        )
        .unwrap();
        assert!(quote.amount_out <= 500);
        assert_eq!(quote.amount_out, 500);
        assert_eq!(quote.amount_in, 500);
    }

    #[rstest]
    #[case(SwapMode::ExactIn, false, Err(PARTIAL_FILL_NOT_ALLOWED))]
    #[case(SwapMode::ExactIn, true, Ok(()))]
    #[case(SwapMode::ExactOut, false, Err(PARTIAL_FILL_NOT_ALLOWED))]
    #[case(SwapMode::ExactOut, true, Ok(()))]
    fn test_swap_quote_partial_fill(
        #[case] swap_mode: SwapMode,
        #[case] allow_partial_fill: bool,
        #[case] expected: Result<(), CoreError>,
    ) {
        // Input upper 10_000, balance 9_000 → room 1_000; request 5_000
        let input = SwapVault {
            swap_fee_per_m: 0,
            inventory_limit_upper: 10_000,
            inventory_limit_lower: 0,
            max_swap_amount: u64::MAX,
            positive_skew_per_m: 0,
            negative_skew_per_m: 0,
            lower_skew_offset: 0,
            upper_skew_offset: 10_000,
            skew_exponent: SkewExponent::Linear,
        };
        let output = SwapVault {
            swap_fee_per_m: 0,
            inventory_limit_upper: 100_000,
            inventory_limit_lower: 0,
            max_swap_amount: u64::MAX,
            positive_skew_per_m: 0,
            negative_skew_per_m: 0,
            lower_skew_offset: 0,
            upper_skew_offset: 100_000,
            skew_exponent: SkewExponent::Linear,
        };
        let result = swap_quote(
            5_000,
            swap_mode,
            U128::from(1u128 << 64),
            allow_partial_fill,
            input,
            9_000,
            output,
            1_000,
        );
        match expected {
            Ok(()) => {
                let quote = result.unwrap();
                assert!(quote.amount_in < 5_000 || quote.amount_out < 5_000);
            }
            Err(e) => assert_eq!(result, Err(e)),
        }
    }

    #[rstest]
    #[case(SwapMode::ExactIn, 1_000)]
    #[case(SwapMode::ExactOut, 1_000)]
    fn test_swap_quote_multiband(#[case] swap_mode: SwapMode, #[case] amount: u64) {
        // Start just into band +1; 1000 crosses into band +2 (width 500)
        let input = skew_vault(32_000, 0, SkewExponent::Linear);
        let output = dead_zone_vault(0, u64::MAX);
        let small = swap_quote(
            100,
            swap_mode,
            U128::from(1u128 << 64),
            true,
            input,
            48_001,
            output,
            50_000,
        )
        .unwrap();
        let large = swap_quote(
            amount,
            swap_mode,
            U128::from(1u128 << 64),
            true,
            input,
            48_001,
            output,
            50_000,
        )
        .unwrap();
        assert!(large.amount_in > 0 && large.amount_out > 0);
        assert!(large.amount_in >= small.amount_in);
        // Crossed into higher positive-skew bands → average out/in rises vs single-band fill
        let small_rate = small.amount_out as u128 * 1_000_000 / small.amount_in.max(1) as u128;
        let large_rate = large.amount_out as u128 * 1_000_000 / large.amount_in.max(1) as u128;
        assert!(large_rate >= small_rate);
    }

    #[test]
    fn test_swap_quote_capacity_overflow_input_above_upper() {
        // balance > upper → checked_sub overflows
        let vault = dead_zone_vault(0, u64::MAX);
        let result = swap_quote(
            1_000,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            vault,
            100_001,
            vault,
            50_000,
        );
        assert_eq!(result, Err(ARITHMETIC_OVERFLOW));
    }

    #[test]
    fn test_swap_quote_capacity_overflow_output_below_lower() {
        let input = dead_zone_vault(0, u64::MAX);
        let output = SwapVault {
            swap_fee_per_m: 0,
            inventory_limit_upper: 100_000,
            inventory_limit_lower: 1_000,
            max_swap_amount: u64::MAX,
            positive_skew_per_m: 0,
            negative_skew_per_m: 0,
            lower_skew_offset: 1_000,
            upper_skew_offset: 100_000,
            skew_exponent: SkewExponent::Linear,
        };
        let result = swap_quote(
            1_000,
            SwapMode::ExactOut,
            U128::from(1u128 << 64),
            true,
            input,
            50_000,
            output,
            500,
        );
        assert_eq!(result, Err(ARITHMETIC_OVERFLOW));
    }

    #[test]
    fn test_swap_quote_zero_prices_exit_loop() {
        // Both-toward at the outer bands with 100% intensity clamps the factor to 0.
        let vault = skew_vault(0, 1_000_000, SkewExponent::Linear);
        let quote = swap_quote(
            400,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            vault,
            0,
            vault,
            64_000,
        )
        .unwrap();
        assert_eq!(quote.amount_in, 0);
        assert_eq!(quote.amount_out, 0);
        assert_eq!(
            swap_quote(
                400,
                SwapMode::ExactIn,
                U128::from(1u128 << 64),
                false,
                vault,
                0,
                vault,
                64_000,
            ),
            Err(PARTIAL_FILL_NOT_ALLOWED)
        );
    }

    #[test]
    fn test_swap_quote_both_away_stacks() {
        let input = skew_vault(32_000, 0, SkewExponent::Linear);
        let output = skew_vault(32_000, 0, SkewExponent::Linear);
        let one_away = swap_quote(
            400,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            input,
            55_501,
            dead_zone_vault(0, u64::MAX),
            32_000,
        )
        .unwrap();
        let both_away = swap_quote(
            400,
            SwapMode::ExactIn,
            U128::from(1u128 << 64),
            true,
            input,
            55_501,
            output,
            8_000,
        )
        .unwrap();
        assert_eq!(one_away.amount_out, 406);
        assert!(both_away.amount_out > one_away.amount_out);
    }
}