riptide-amm-math 2.0.1

use borsh::{BorshDeserialize, BorshSerialize};

use ethnum::U256;

use super::{
    super::error::{
        CoreError, AMOUNT_EXCEEDS_MAX_U128, AMOUNT_EXCEEDS_MAX_U64, ARITHMETIC_OVERFLOW,
        INVALID_ORACLE_DATA,
    },
    QuoteType, SingleSideLiquidity, LIQUIDITY_LEVELS, PER_M_DENOMINATOR,
};

const MIN_SQRT_PRICE: u128 = 5647135299341; // ~1e-13
const MAX_SQRT_PRICE: u128 = 60257519765924248467716150; // ~2e13

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

const MAX_POSITIONS: usize = 16;

// Minimum sqrt price range width to ensure virtual reserves are non-zero.
// Derived from: L * (upper - lower) >> 64 >= 1, where L = PER_M_DENOMINATOR (1_000_000).
// So (upper - lower) >= ceil(2^64 / 1_000_000) = 18_446_744_073_710.
const MIN_SQRT_PRICE_RANGE: u128 = 18_446_744_073_710;

// A liquidity position here works slightly differently than a regular AMM.
// Since the liquidity is discretized into 32 levels, we need to calculate the price for each level.
// A normal AMM would have a continuous price function across the range.
// A position here is set up in such a way that the first price is always
// the mid price (sans spread) and the last price is always the outer edge of the position's price
// range.

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[cfg_attr(true, derive(BorshDeserialize, BorshSerialize))]
#[cfg_attr(feature = "wasm", wasm_expose)]
pub enum LiquidityType {
    // CP shouldn't be used since it is wildly inaccurate if spread across 32 liquidity levels
    ConstantProduct,
    ConcentratedLiquidity {
        lower_sqrt_price: u128,
        upper_sqrt_price: u128,
    },
}

impl LiquidityType {
    fn positions(&self) -> Result<[Position; MAX_POSITIONS], CoreError> {
        let mut positions = [Position::default(); MAX_POSITIONS];
        match self {
            LiquidityType::ConstantProduct => {
                positions[0] = Position::full_range(PER_M_DENOMINATOR as u32);
            }
            LiquidityType::ConcentratedLiquidity {
                lower_sqrt_price,
                upper_sqrt_price,
            } => {
                positions[0] = Position::new(
                    *lower_sqrt_price,
                    *upper_sqrt_price,
                    PER_M_DENOMINATOR as u32,
                )?;
            }
        };
        Ok(positions)
    }
}

#[derive(Debug, Clone, Copy)]
struct Position {
    lower_sqrt_price: u128,
    upper_sqrt_price: u128,
    liquidity_share_per_m: u32,
}

impl Default for Position {
    fn default() -> Self {
        Self {
            lower_sqrt_price: MIN_SQRT_PRICE,
            upper_sqrt_price: MAX_SQRT_PRICE,
            liquidity_share_per_m: 0,
        }
    }
}

impl Position {
    fn full_range(liquidity_share_per_m: u32) -> Self {
        Self {
            lower_sqrt_price: MIN_SQRT_PRICE,
            upper_sqrt_price: MAX_SQRT_PRICE,
            liquidity_share_per_m,
        }
    }

    fn new(
        lower_sqrt_price: u128,
        upper_sqrt_price: u128,
        liquidity_share_per_m: u32,
    ) -> Result<Self, CoreError> {
        if lower_sqrt_price > upper_sqrt_price
            || lower_sqrt_price < MIN_SQRT_PRICE
            || upper_sqrt_price > MAX_SQRT_PRICE
            || upper_sqrt_price - lower_sqrt_price < MIN_SQRT_PRICE_RANGE
        {
            return Err(INVALID_ORACLE_DATA);
        }
        Ok(Self {
            lower_sqrt_price,
            upper_sqrt_price,
            liquidity_share_per_m,
        })
    }
}

// AMM oracle: current price P is implied by reserves and positions (no stored price).
// We solve for P using the side being built only: bid => Y(s) = reserves_b, ask => X(s) =
// reserves_a. Then L_total is scaled so that virtual reserve matches the reserve for that side.

pub(crate) fn amm_price(
    liquidity_type: LiquidityType,
    reserves_a: u64,
    reserves_b: u64,
) -> Result<u128, CoreError> {
    if reserves_a == 0 && reserves_b == 0 {
        return Ok(0);
    }

    let positions = liquidity_type.positions()?;
    let sqrt_price = implied_sqrt_price(&positions, reserves_a, reserves_b)?;

    U256::from(sqrt_price)
        .checked_mul(U256::from(sqrt_price))
        .ok_or(ARITHMETIC_OVERFLOW)?
        .checked_shr(64)
        .ok_or(ARITHMETIC_OVERFLOW)?
        .try_into()
        .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)
}

pub(crate) fn amm_liquidity(
    liquidity_type: LiquidityType,
    bid_spread_per_m: i32,
    ask_spread_per_m: i32,
    quote_type: QuoteType,
    reserves_a: u64,
    reserves_b: u64,
) -> Result<SingleSideLiquidity, CoreError> {
    if bid_spread_per_m <= -PER_M_DENOMINATOR
        || bid_spread_per_m >= PER_M_DENOMINATOR
        || ask_spread_per_m <= -PER_M_DENOMINATOR
    {
        return Err(INVALID_ORACLE_DATA);
    }
    if reserves_a == 0 && reserves_b == 0 {
        return Ok(SingleSideLiquidity::new());
    }

    let positions = liquidity_type.positions()?;
    let total_liquidity_share_per_m = positions
        .iter()
        .map(|p| u128::from(p.liquidity_share_per_m))
        .sum::<u128>();
    if total_liquidity_share_per_m == 0 || total_liquidity_share_per_m > PER_M_DENOMINATOR as u128 {
        return Err(INVALID_ORACLE_DATA);
    }

    let mid_sqrt_price = implied_sqrt_price(&positions, reserves_a, reserves_b)?;

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

    let spread = if quote_type.a_to_b() {
        PER_M_DENOMINATOR - bid_spread_per_m
    } else {
        PER_M_DENOMINATOR + ask_spread_per_m
    };

    let end_sqrt_price = if quote_type.a_to_b() {
        positions
            .iter()
            .filter(|p| p.liquidity_share_per_m > 0)
            .map(|p| p.lower_sqrt_price)
            .min()
            .unwrap_or(MIN_SQRT_PRICE)
    } else {
        positions
            .iter()
            .filter(|p| p.liquidity_share_per_m > 0)
            .map(|p| p.upper_sqrt_price)
            .max()
            .unwrap_or(MAX_SQRT_PRICE)
    };

    let mut liquidity_bins: [(u128, u64); LIQUIDITY_LEVELS] = [(0, 0); LIQUIDITY_LEVELS];
    let mut position_bin_indexes: [[bool; LIQUIDITY_LEVELS]; MAX_POSITIONS] =
        [[false; LIQUIDITY_LEVELS]; MAX_POSITIONS];

    // Calculate the price for each liquidity level
    for i in 0..LIQUIDITY_LEVELS {
        let (start_sqrt_price, end_sqrt_price) = bin_boundaries(i, mid_sqrt_price, end_sqrt_price)?;

        // For each position, mark the bin as contributing whenever there is
        // ANY non-zero overlap between the bin's sqrt range and the position's
        // range — not only when the bin is fully contained. A partial overlap
        // should still receive its proportional share of the position's
        // reserves; treating it as zero would dilute liquidity at every bin
        // that straddles a position boundary (see audit P2-L-02).
        for j in 0..MAX_POSITIONS {
            position_bin_indexes[j][i] =
                bin_position_overlap(i, mid_sqrt_price, end_sqrt_price, &positions[j])? > 0;
        }

        liquidity_bins[i].0 = bin_price(i, start_sqrt_price, end_sqrt_price, spread)?;
    }

    let position_amounts =
        position_reserves(&positions, mid_sqrt_price, reserves, quote_type.a_to_b())?;

    // Add the amounts to the liquidity levels proportional to overlap width.
    // Two-pass approach to avoid storing a separate overlap_widths array on the stack.
    for i in 0..MAX_POSITIONS {
        let num_bins = position_bin_indexes[i].iter().filter(|&b| *b).count();
        if num_bins == 0 {
            continue;
        }
        if positions[i].liquidity_share_per_m == 0 {
            continue;
        }

        // Pass 1: compute total overlap width
        let mut total_overlap: u128 = 0;
        for j in 0..LIQUIDITY_LEVELS {
            if position_bin_indexes[i][j] {
                total_overlap = total_overlap
                    .checked_add(bin_position_overlap(
                        j,
                        mid_sqrt_price,
                        end_sqrt_price,
                        &positions[i],
                    )?)
                    .ok_or(ARITHMETIC_OVERFLOW)?;
            }
        }

        if total_overlap == 0 {
            continue;
        }

        // Pass 2: distribute proportionally
        for j in 0..LIQUIDITY_LEVELS {
            if position_bin_indexes[i][j] {
                let width = bin_position_overlap(j, mid_sqrt_price, end_sqrt_price, &positions[i])?;
                if width > 0 {
                    let amount = U256::from(position_amounts[i])
                        .checked_mul(U256::from(width))
                        .ok_or(ARITHMETIC_OVERFLOW)?
                        .checked_div(U256::from(total_overlap))
                        .ok_or(ARITHMETIC_OVERFLOW)?;
                    let amount: u64 = amount.try_into().map_err(|_| AMOUNT_EXCEEDS_MAX_U64)?;
                    liquidity_bins[j].1 = liquidity_bins[j]
                        .1
                        .checked_add(amount)
                        .ok_or(ARITHMETIC_OVERFLOW)?;
                }
            }
        }
    }

    // Normalize the total liquidity to reserves
    let total_amount = liquidity_bins
        .iter()
        .map(|b: &(u128, u64)| b.1)
        .sum::<u64>();

    let mut remaining = reserves.saturating_sub(total_amount);
    // If every bin floored to zero, there is nothing to distribute the
    // residual onto — an entry-only loop would spin forever. Skip the
    // normalization in that case; the unallocated residual is at most
    // a few units (bounded by per-bin floor errors).
    let any_bin_has_liquidity = liquidity_bins.iter().any(|b| b.1 > 0);
    if any_bin_has_liquidity {
        while remaining > 0 {
            for i in (0..LIQUIDITY_LEVELS).rev() {
                if remaining == 0 {
                    break;
                }
                if liquidity_bins[i].1 > 0 {
                    liquidity_bins[i].1 = liquidity_bins[i]
                        .1
                        .checked_add(1)
                        .ok_or(ARITHMETIC_OVERFLOW)?;
                    remaining -= 1;
                }
            }
        }
    }

    Ok(SingleSideLiquidity::from_slice(&liquidity_bins))
}

fn virtual_position_reserves(
    sqrt_price: u128,
    position: &Position,
) -> Result<(u64, u64), CoreError> {
    let liquidity = U256::from(position.liquidity_share_per_m);
    let sqrt_price = sqrt_price.clamp(position.lower_sqrt_price, position.upper_sqrt_price);

    let diff_a = U256::from(position.upper_sqrt_price).saturating_sub(U256::from(sqrt_price));
    let virtual_a = if diff_a == 0 {
        0u64
    } else {
        let numerator = liquidity
            .checked_mul(diff_a)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .checked_shl(64)
            .ok_or(ARITHMETIC_OVERFLOW)?;
        let denominator = U256::from(sqrt_price)
            .checked_mul(U256::from(position.upper_sqrt_price))
            .ok_or(ARITHMETIC_OVERFLOW)?;

        numerator
            .checked_div(denominator)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U64)?
    };

    let diff_b = U256::from(sqrt_price).saturating_sub(U256::from(position.lower_sqrt_price));
    let virtual_b = if diff_b == 0 {
        0u64
    } else {
        liquidity
            .checked_mul(diff_b)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .checked_shr(64)
            .ok_or(ARITHMETIC_OVERFLOW)?
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U64)?
    };

    Ok((virtual_a, virtual_b))
}

fn virtual_reserves(positions: &[Position], sqrt_price: u128) -> Result<(u64, u64), CoreError> {
    let mut virtual_a = 0u64;
    let mut virtual_b = 0u64;
    for position in positions {
        if position.liquidity_share_per_m == 0 {
            continue;
        }
        let (a_amount, b_amount) = virtual_position_reserves(sqrt_price, position)?;
        virtual_a = virtual_a.checked_add(a_amount).ok_or(ARITHMETIC_OVERFLOW)?;
        virtual_b = virtual_b.checked_add(b_amount).ok_or(ARITHMETIC_OVERFLOW)?;
    }
    Ok((virtual_a, virtual_b))
}

fn position_reserves(
    positions: &[Position],
    sqrt_price: u128,
    reserves: u64,
    a_to_b: bool,
) -> Result<[u64; MAX_POSITIONS], CoreError> {
    // Calculate the virtual amounts for each position
    let mut total_virtual_amount: u64 = 0;
    let mut virtual_amounts: [u64; MAX_POSITIONS] = [0; MAX_POSITIONS];
    for i in 0..MAX_POSITIONS {
        let (virtual_a, virtual_b) = virtual_position_reserves(sqrt_price, &positions[i])?;
        if a_to_b {
            virtual_amounts[i] = virtual_b;
            total_virtual_amount = total_virtual_amount
                .checked_add(virtual_b)
                .ok_or(ARITHMETIC_OVERFLOW)?;
        } else {
            virtual_amounts[i] = virtual_a;
            total_virtual_amount = total_virtual_amount
                .checked_add(virtual_a)
                .ok_or(ARITHMETIC_OVERFLOW)?;
        }
    }

    // Calculate the actual amounts for each position
    let mut actual_amounts: [u64; MAX_POSITIONS] = [0; MAX_POSITIONS];
    for i in 0..MAX_POSITIONS {
        actual_amounts[i] = reserves
            .checked_mul(virtual_amounts[i])
            .ok_or(ARITHMETIC_OVERFLOW)?
            .checked_div(total_virtual_amount)
            .ok_or(ARITHMETIC_OVERFLOW)?;
    }

    // Normalize sum(actual_a_amounts) to reserves_a so that we're using all of the reserves
    let total_actual_amount = actual_amounts.iter().sum::<u64>();
    let mut remaining = reserves.saturating_sub(total_actual_amount);
    // Same guard as in `amm_liquidity`: if every position floored to zero
    // the inner check would never fire and the outer loop would spin
    // forever. Skip the normalization; the residual is bounded.
    let any_position_has_amount = actual_amounts.iter().any(|&a| a > 0);
    if any_position_has_amount {
        while remaining > 0 {
            for i in 0..MAX_POSITIONS {
                if remaining == 0 {
                    break;
                }
                if actual_amounts[i] > 0 {
                    actual_amounts[i] = actual_amounts[i]
                        .checked_add(1)
                        .ok_or(ARITHMETIC_OVERFLOW)?;
                    remaining -= 1;
                }
            }
        }
    }

    Ok(actual_amounts)
}

fn implied_sqrt_price(
    positions: &[Position],
    reserves_a: u64,
    reserves_b: u64,
) -> Result<u128, CoreError> {
    if positions.is_empty() {
        return Err(INVALID_ORACLE_DATA);
    }

    let mut lowest_sqrt_price = positions
        .iter()
        .filter(|p| p.liquidity_share_per_m > 0)
        .map(|p| p.lower_sqrt_price)
        .min()
        .unwrap_or(0);
    let mut highest_sqrt_price = positions
        .iter()
        .filter(|p| p.liquidity_share_per_m > 0)
        .map(|p| p.upper_sqrt_price)
        .max()
        .unwrap_or(u128::MAX);

    // If one reserve is zero, the sqrt price is the highest or lowest price.
    if reserves_a == 0 {
        return Ok(highest_sqrt_price);
    } else if reserves_b == 0 {
        return Ok(lowest_sqrt_price);
    }

    // Binary search for P where virtual_b(P)/virtual_a(P) = reserves_b/reserves_a.
    for _ in 0..128 {
        if highest_sqrt_price <= lowest_sqrt_price + 1 {
            break;
        }
        let diff = highest_sqrt_price.abs_diff(lowest_sqrt_price) >> 1;
        let mid = lowest_sqrt_price
            .checked_add(diff)
            .ok_or(ARITHMETIC_OVERFLOW)?;
        let (virtual_a, virtual_b) = virtual_reserves(positions, mid)?;

        let lhs = u128::from(virtual_b)
            .checked_mul(reserves_a as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?;
        let rhs = u128::from(virtual_a)
            .checked_mul(reserves_b as u128)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        if lhs < rhs {
            lowest_sqrt_price = mid;
        } else {
            highest_sqrt_price = mid;
        }
    }

    let diff = highest_sqrt_price.abs_diff(lowest_sqrt_price) >> 1;
    let sqrt_price = lowest_sqrt_price
        .checked_add(diff)
        .ok_or(ARITHMETIC_OVERFLOW)?;

    Ok(sqrt_price)
}

#[inline(always)]
fn bin_position_overlap(
    bin_index: usize,
    mid_sqrt_price: u128,
    end_sqrt_price: u128,
    position: &Position,
) -> Result<u128, CoreError> {
    let (bin_start, bin_end) = bin_boundaries(bin_index, mid_sqrt_price, end_sqrt_price)?;
    let (lo, hi) = if bin_start <= bin_end {
        (bin_start, bin_end)
    } else {
        (bin_end, bin_start)
    };
    let overlap_start = lo.max(position.lower_sqrt_price);
    let overlap_end = hi.min(position.upper_sqrt_price);
    if overlap_end > overlap_start {
        Ok(overlap_end - overlap_start)
    } else {
        Ok(0)
    }
}

fn bin_boundaries(
    i: usize,
    sqrt_price: u128,
    end_sqrt_price: u128,
) -> Result<(u128, u128), CoreError> {
    let sqrt_price_diff = end_sqrt_price.abs_diff(sqrt_price);

    let start_sqrt_price_delta = U256::from(sqrt_price_diff)
        .checked_mul(U256::from(i as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?
        .checked_div(U256::from(LIQUIDITY_LEVELS as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let end_sqrt_price_delta = U256::from(sqrt_price_diff)
        .checked_mul(U256::from(i as u128 + 1))
        .ok_or(ARITHMETIC_OVERFLOW)?
        .checked_div(U256::from(LIQUIDITY_LEVELS as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let (start_sqrt_price, end_sqrt_price) = if end_sqrt_price > sqrt_price {
        let start = U256::from(sqrt_price)
            .checked_add(start_sqrt_price_delta)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        let end = U256::from(sqrt_price)
            .checked_add(end_sqrt_price_delta)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        (start, end)
    } else {
        let start = U256::from(sqrt_price)
            .checked_sub(start_sqrt_price_delta)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        let end = U256::from(sqrt_price)
            .checked_sub(end_sqrt_price_delta)
            .ok_or(ARITHMETIC_OVERFLOW)?;

        (start, end)
    };

    Ok((
        start_sqrt_price
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?,
        end_sqrt_price
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?,
    ))
}

// Take the first price in the bin and apply the spread. If we wanted to be more true to an actual
// AMM, we would take the geometric mean. The problem if we do that is that the first bin's price
// would not be equal to the mid price meaning you would get an implicit spread which becomes larger
// the larger the lp range.
fn bin_price(
    i: usize,
    start_sqrt_price: u128,
    end_sqrt_price: u128,
    spread: i32,
) -> Result<u128, CoreError> {
    let sqrt_price_diff = end_sqrt_price.abs_diff(start_sqrt_price);
    let sqrt_price_delta = U256::from(sqrt_price_diff)
        .checked_mul(U256::from(i as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?
        .checked_div(U256::from(LIQUIDITY_LEVELS as u128 - 1))
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let sqrt_price = if end_sqrt_price > start_sqrt_price {
        U256::from(start_sqrt_price)
            .checked_add(sqrt_price_delta)
            .ok_or(ARITHMETIC_OVERFLOW)?
    } else {
        U256::from(start_sqrt_price)
            .checked_sub(sqrt_price_delta)
            .ok_or(ARITHMETIC_OVERFLOW)?
    };

    let product = sqrt_price
        .checked_mul(sqrt_price)
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let quotient = product.checked_shr(64).ok_or(ARITHMETIC_OVERFLOW)?;
    // Only the low 64 bits were discarded by the right shift, so the remainder of the
    // division by 2^64 is bits[0..64]. Masking with u128::MAX previously bled the low
    // half of the quotient into the remainder check and forced a spurious round-up
    // whenever bits[64..128] of the product were set (e.g. sqrt_price = 2^44).
    let remainder = product & U256::from(u64::MAX);

    let base_price: u128 = if end_sqrt_price > start_sqrt_price && remainder > 0 {
        (quotient + 1)
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?
    } else {
        quotient.try_into().map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?
    };

    let product = U256::from(base_price)
        .checked_mul(U256::from(spread as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let quotient = product
        .checked_div(U256::from(PER_M_DENOMINATOR as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?;
    let remainder = product
        .checked_rem(U256::from(PER_M_DENOMINATOR as u128))
        .ok_or(ARITHMETIC_OVERFLOW)?;

    let price = if end_sqrt_price > start_sqrt_price && remainder > 0 {
        (quotient + 1)
            .try_into()
            .map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?
    } else {
        quotient.try_into().map_err(|_| AMOUNT_EXCEEDS_MAX_U128)?
    };

    Ok(price)
}

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

    #[rstest]
    #[case(Position::full_range(PER_M_DENOMINATOR as u32), 100, 100, 18446738426575981041)]
    #[case(Position::full_range(PER_M_DENOMINATOR as u32), 150, 50, 10650233338091508966)]
    #[case(Position::full_range(PER_M_DENOMINATOR as u32), 50, 150, 31950688720003928217)]
    #[case(Position::full_range(PER_M_DENOMINATOR as u32), 200, 0, MIN_SQRT_PRICE)]
    #[case(Position::full_range(PER_M_DENOMINATOR as u32), 0, 200, MAX_SQRT_PRICE)]
    #[case(Position::new((1 << 64) / 2, (1 << 64) * 2, PER_M_DENOMINATOR as u32).unwrap(), 100, 100, 18446744073709551615)]
    #[case(Position::new((1 << 64) / 3, (1 << 64) * 2, PER_M_DENOMINATOR as u32).unwrap(), 100, 100, 16973448724429105277)]
    #[case(Position::new((1 << 64) / 2, (1 << 64) * 3, PER_M_DENOMINATOR as u32).unwrap(), 100, 100, 20047903282541897858)]
    #[case(Position::new((1 << 64) / 2, (1 << 64) * 2, PER_M_DENOMINATOR as u32).unwrap(), 150, 50, 14159573177026862144)]
    #[case(Position::new((1 << 64) / 2, (1 << 64) * 2, PER_M_DENOMINATOR as u32).unwrap(), 50, 150, 24031964994045732128)]
    #[case(Position::new((1 << 64) / 2, (1 << 64) * 2, PER_M_DENOMINATOR as u32).unwrap(), 200, 0, (1 << 64) / 2)]
    #[case(Position::new((1 << 64) / 2, (1 << 64) * 2, PER_M_DENOMINATOR as u32).unwrap(), 0, 200, (1 << 64) * 2)]
    fn test_implied_sqrt_price_single_position(
        #[case] position: Position,
        #[case] reserves_a: u64,
        #[case] reserves_b: u64,
        #[case] expected: u128,
    ) {
        let sqrt_price = implied_sqrt_price(&[position], reserves_a, reserves_b).unwrap();
        assert_eq!(sqrt_price, expected);
    }

    #[rstest]
    #[case(100, 100, 18446744073709551615)]
    #[case(150, 50, 13516113850247725564)]
    #[case(50, 150, 25176050652658693911)]
    #[case(200, 0, 6148914691236517205)]
    #[case(0, 200, 55340232221128654848)]
    fn test_implied_sqrt_price_multiple_positions(
        #[case] reserves_a: u64,
        #[case] reserves_b: u64,
        #[case] expected: u128,
    ) {
        let positions = vec![
            Position::new((1 << 64) / 2, (1 << 64) * 2, 300_000).unwrap(),
            Position::new((1 << 64) / 2, (1 << 64) * 3, 200_000).unwrap(),
            Position::new((1 << 64) / 3, (1 << 64) * 2, 200_000).unwrap(),
            Position::new((1 << 64) / 3, (1 << 64) * 3, 300_000).unwrap(),
        ];
        let sqrt_price = implied_sqrt_price(&positions, reserves_a, reserves_b).unwrap();
        assert_eq!(sqrt_price, expected);
    }

    #[rstest]
    #[case(LiquidityType::ConstantProduct, 100, 100, Ok(18446732779444139232))]
    #[case(LiquidityType::ConstantProduct, 150, 50, Ok(6148915478122426543))]
    #[case(LiquidityType::ConstantProduct, 50, 150, Ok(55340200178605227858))]
    #[case(LiquidityType::ConstantProduct, 200, 0, Ok(1728767))]
    #[case(
        LiquidityType::ConstantProduct,
        0,
        200,
        Ok(196835207006294262292126795817913)
    )]
    #[case(LiquidityType::ConcentratedLiquidity { lower_sqrt_price: (1 << 64) / 2, upper_sqrt_price: (1 << 64) * 2 }, 100, 100, Ok(18446744073709551614))]
    #[case(LiquidityType::ConcentratedLiquidity { lower_sqrt_price: (1 << 64) / 2, upper_sqrt_price: (1 << 64) * 2 }, 150, 50, Ok(10868775094100403166))]
    #[case(LiquidityType::ConcentratedLiquidity { lower_sqrt_price: (1 << 64) / 2, upper_sqrt_price: (1 << 64) * 2 }, 50, 150, Ok(31308253595719773170))]
    #[case(LiquidityType::ConcentratedLiquidity { lower_sqrt_price: (1 << 64) / 2, upper_sqrt_price: (1 << 64) * 2 }, 200, 0, Ok(4611686018427387904))]
    #[case(LiquidityType::ConcentratedLiquidity { lower_sqrt_price: (1 << 64) / 2, upper_sqrt_price: (1 << 64) * 2 }, 0, 200, Ok(73786976294838206464))]
    fn test_amm_price(
        #[case] liquidity_type: LiquidityType,
        #[case] reserves_a: u64,
        #[case] reserves_b: u64,
        #[case] expected: Result<u128, CoreError>,
    ) {
        let price = amm_price(liquidity_type, reserves_a, reserves_b);
        assert_eq!(price, expected);
    }

    #[rstest]
    #[case(0, MIN_SQRT_PRICE, Ok((18446744073709551616, 17870283497879106233)))]
    #[case(1, MIN_SQRT_PRICE, Ok((17870283497879106233, 17293822922048660849)))]
    #[case(0, MAX_SQRT_PRICE, Ok((18446744073709551616, 1883065362968454170744257)))]
    #[case(1, MAX_SQRT_PRICE, Ok((1883065362968454170744257, 3766112279192834631936899)))]
    fn test_bin_boundaries(
        #[case] i: usize,
        #[case] end_sqrt_price: u128,
        #[case] expected: Result<(u128, u128), CoreError>,
    ) {
        let result = bin_boundaries(i, 1 << 64, end_sqrt_price);
        assert_eq!(result, expected);
    }

    #[rstest]
    #[case(0, 1 << 64, (1 << 64) / 2, 999_000, Ok(18428297329635842064))]
    #[case(31, 1 << 64, (1 << 64) / 2, 999_000, Ok(4607074332408960516))]
    #[case(0, 1 << 64, (1 << 64) / 2, 998_000, Ok(18409850585562132512))]
    #[case(31, 1 << 64, (1 << 64) / 2, 998_000, Ok(4602462646390533128))]
    #[case(0, 1 << 64, (1 << 64) / 2, 1_001_000, Ok(18465190817783261167))]
    #[case(31, 1 << 64, (1 << 64) / 2, 1_002_000, Ok(4620909390464242679))]
    #[case(0, 1 << 64, (1 << 64) * 2, 999_000, Ok(18428297329635842065))]
    #[case(31, 1 << 64, (1 << 64) * 2, 999_000, Ok(73713189318543368258))]
    #[case(0, 1 << 64, (1 << 64) * 2, 998_000, Ok(18409850585562132513))]
    #[case(31, 1 << 64, (1 << 64) * 2, 998_000, Ok(73639402342248530052))]
    #[case(0, 1 << 64, (1 << 64) * 2, 1_001_000, Ok(18465190817783261168))]
    #[case(31, 1 << 64, (1 << 64) * 2, 1_001_000, Ok(73860763271133044671))]
    #[case(0, 1 << 64, (1 << 64) * 2, 1_002_000, Ok(18483637561856970720))]
    #[case(31, 1 << 64, (1 << 64) * 2, 1_002_000, Ok(73934550247427882877))]
    fn test_bin_price(
        #[case] i: usize,
        #[case] start_sqrt_price: u128,
        #[case] end_sqrt_price: u128,
        #[case] spread: i32,
        #[case] expected: Result<u128, CoreError>,
    ) {
        let result = bin_price(i, start_sqrt_price, end_sqrt_price, spread);
        assert_eq!(result, expected);
    }

    // Q64.64 squaring is exact when the sqrt price is a pure power of two whose square
    // is at least 2^64 (e.g. sqrt_price = 2^44 -> product = 2^88, quotient = 2^24,
    // remainder = 0). The previous remainder mask (`u128::MAX`) bled the low 64 bits of
    // the quotient back into the remainder check and forced a spurious +1 round-up on
    // increasing bins. With i=0 the sqrt-price delta is zero, so the result reduces to
    // the squared start_sqrt_price; combined with spread = PER_M_DENOMINATOR the second
    // division is also exact, leaving any +1 here attributable solely to the masking bug.
    #[rstest]
    #[case(1u128 << 44, 1 << 24)]
    #[case(1u128 << 50, 1 << 36)]
    #[case(1u128 << 60, 1 << 56)]
    fn test_bin_price_no_spurious_round_up_on_exact_division(
        #[case] start_sqrt_price: u128,
        #[case] expected: u128,
    ) {
        let result = bin_price(
            0,
            start_sqrt_price,
            start_sqrt_price + 1, // increasing direction so the round-up branch is reachable
            PER_M_DENOMINATOR,    // spread of 1.0 -> second division is exact
        );
        assert_eq!(result, Ok(expected));
    }

    #[test]
    fn test_amm_liquidity_a_to_b() {
        let liquidity = amm_liquidity(
            LiquidityType::ConstantProduct,
            10000,
            20000,
            QuoteType::TokenAExactIn,
            1000,
            1000,
        )
        .unwrap()
        .as_slice()
        .to_vec();
        // Prices starts ~ 0.99 and ends ~ 9e-14
        let expected = vec![
            (18262265451649697839, 31),
            (17103058524375092466, 31),
            (15981858369775465113, 31),
            (14898664987850815784, 31),
            (13853478378601144476, 31),
            (12846298542026451190, 31),
            (11877125478126735926, 31),
            (10945959186901998683, 31),
            (10052799668352239462, 31),
            (9197646922477458264, 31),
            (8380500949277655088, 31),
            (7601361748752829935, 31),
            (6860229320902982803, 31),
            (6157103665728113694, 31),
            (5491984783228222606, 31),
            (4864872673403309539, 31),
            (4275767336253374494, 31),
            (3724668771778417472, 31),
            (3211576979978438473, 31),
            (2736491960853437495, 31),
            (2299413714403414540, 31),
            (1900342240628369606, 31),
            (1539277539528302694, 31),
            (1216219611103213804, 31),
            (931168455353102936, 32),
            (684124072277970090, 32),
            (475086461877815267, 32),
            (304055624152638466, 32),
            (171031559102439686, 32),
            (76014266727218928, 32),
            (19003747026976192, 32),
            (1711479, 32),
        ];

        assert_eq!(liquidity, expected);
    }

    #[test]
    fn test_amm_liquidity_concentrated_a_to_b() {
        let liquidity = amm_liquidity(
            LiquidityType::ConcentratedLiquidity {
                lower_sqrt_price: (1 << 64) / 2,
                upper_sqrt_price: (1 << 64) * 2,
            },
            10000,
            20000,
            QuoteType::TokenAExactIn,
            1000,
            1000,
        )
        .unwrap()
        .as_slice()
        .to_vec();
        // Prices starts ~ 0.99 and ends ~ 0.2475
        let expected = vec![
            (18262276632972456097, 31),
            (17677921787536552847, 31),
            (17103068646904485422, 31),
            (16537717211076253820, 31),
            (15981867480051858042, 31),
            (15435519453831298092, 31),
            (14898673132414573967, 31),
            (14371328515801685667, 31),
            (13853485603992633191, 31),
            (13345144396987416541, 31),
            (12846304894786035717, 31),
            (12356967097388490717, 31),
            (11877131004794781542, 31),
            (11406796617004908193, 31),
            (10945963934018870670, 31),
            (10494632955836668971, 31),
            (10052803682458303097, 31),
            (9620476113883773049, 31),
            (9197650250113078826, 31),
            (8784326091146220429, 31),
            (8380503636983197855, 31),
            (7986182887624011109, 31),
            (7601363843068660187, 31),
            (7226046503317145091, 31),
            (6860230868369465819, 32),
            (6503916938225622372, 32),
            (6157104712885614751, 32),
            (5819794192349442955, 32),
            (5491985376617106984, 32),
            (5173678265688606840, 32),
            (4864872859563942520, 32),
            (4565569158243114024, 32),
        ];

        assert_eq!(liquidity, expected);
    }

    #[test]
    fn test_amm_liquidity_b_to_a() {
        let liquidity = amm_liquidity(
            LiquidityType::ConstantProduct,
            10000,
            20000,
            QuoteType::TokenBExactIn,
            1000,
            1000,
        )
        .unwrap()
        .as_slice()
        .to_vec();
        // Prices starts ~ 1.02 and ends ~ 1e13
        let expected = vec![
            (18815667435033022018, 31),
            (208923620106592073258895578064, 31),
            (835686549707659367878445172487, 31),
            (1880288788822017551293681805289, 31),
            (3342730337449666623504606336313, 31),
            (5223011195590606584511217260829, 31),
            (7521131363244837434313515223724, 31),
            (10237090840412359172911501729725, 31),
            (13370889627093171800305175704025, 31),
            (16922527723287275316494535211973, 31),
            (20892005128994669721479581758299, 31),
            (25279321844215355015260315343002, 31),
            (30084477868949331197836738545617, 31),
            (35307473203196598269208849216532, 31),
            (40948307846957156229376644346290, 31),
            (47006981800231005078340126514425, 31),
            (53483495063018144816099299160316, 31),
            (60377847635318575442654155620167, 31),
            (67690039517132296958004699118395, 31),
            (75420070708459309362150933739263, 31),
            (83567941209299612655092855828431, 31),
            (92133651019653206836830460871714, 31),
            (101117200139520091907363752953373, 31),
            (110518588568900267866692732073410, 31),
            (120337816307793734714817403390893, 32),
            (130574883356200492451737762176676, 32),
            (141229789714120541077453802841767, 32),
            (152302535381553880591965530545236, 32),
            (163793120358500510995272951305994, 32),
            (175701544644960432287376053301180, 32),
            (188027808240933644468274842334742, 32),
            (200771911146420147537969331734273, 32),
        ];

        assert_eq!(liquidity, expected);
    }

    #[test]
    fn test_amm_liquidity_concentrated_b_to_a() {
        let liquidity = amm_liquidity(
            LiquidityType::ConcentratedLiquidity {
                lower_sqrt_price: (1 << 64) / 2,
                upper_sqrt_price: (1 << 64) * 2,
            },
            10000,
            20000,
            QuoteType::TokenBExactIn,
            1000,
            1000,
        )
        .unwrap()
        .as_slice()
        .to_vec();
        // Prices starts ~ 1.02 and ends ~ 4.08
        let expected = vec![
            (18815678955183742648, 31),
            (20049172996990793413, 31),
            (21321825579807591824, 31),
            (22633636703634137876, 31),
            (23984606368470431575, 31),
            (25374734574316472913, 31),
            (26804021321172261898, 31),
            (28272466609037798524, 31),
            (29780070437913082795, 31),
            (31326832807798114708, 31),
            (32912753718692894266, 31),
            (34537833170597421466, 31),
            (36202071163511696311, 31),
            (37905467697435718797, 31),
            (39648022772369488929, 31),
            (41429736388313006701, 31),
            (43250608545266272120, 31),
            (45110639243229285179, 31),
            (47009828482202045885, 31),
            (48948176262184554231, 31),
            (50925682583176810224, 31),
            (52942347445178813857, 31),
            (54998170848190565136, 31),
            (57093152792212064055, 31),
            (59227293277243310621, 32),
            (61400592303284304828, 32),
            (63613049870335046681, 32),
            (65864665978395536173, 32),
            (68155440627465773314, 32),
            (70485373817545758093, 32),
            (72854465548635490520, 32),
            (75262715820734970594, 32),
        ];

        assert_eq!(liquidity, expected);
    }

    #[test]
    fn test_narrow_range_rejected() {
        // Range width = ~1e-6 in sqrt-price space, too narrow to produce non-zero virtual reserves
        let result = amm_liquidity(
            LiquidityType::ConcentratedLiquidity {
                lower_sqrt_price: 1 << 64,
                upper_sqrt_price: ((1 << 64) * 1_000_001) / 1_000_000,
            },
            0,
            0,
            QuoteType::TokenAExactIn,
            1_000_000,
            1_000_000,
        );
        assert_eq!(result, Err(INVALID_ORACLE_DATA));
    }

    #[test]
    fn test_min_valid_range_accepted() {
        // Range width exactly at MIN_SQRT_PRICE_RANGE should pass Position::new validation
        let result = Position::new(
            1 << 64,
            (1 << 64) + MIN_SQRT_PRICE_RANGE,
            PER_M_DENOMINATOR as u32,
        );
        assert!(result.is_ok());
    }

    // Regression for [P1-L-09] — ensure tiny reserves do not trigger an
    // infinite loop in the reserve-normalization stage. With reserves much
    // smaller than the total virtual amounts, `position_reserves` and the
    // distribution loop in `amm_liquidity` could previously have all entries
    // floor to zero, leaving a non-zero `remaining` with no bin to seed.
    //
    // The function may legitimately return Err for degenerate inputs (e.g.
    // single-side empty); what matters is that the call terminates instead
    // of looping forever. Test harness timeout will catch any regression.
    #[test]
    fn test_tiny_reserves_do_not_loop_forever() {
        let lp = LiquidityType::ConcentratedLiquidity {
            lower_sqrt_price: 1 << 64,
            upper_sqrt_price: (1 << 64) * 2,
        };
        for (ra, rb) in [(1u64, 1), (1, 0), (0, 1), (1, 100), (100, 1)] {
            for qt in [
                QuoteType::TokenAExactIn,
                QuoteType::TokenBExactIn,
                QuoteType::TokenAExactOut,
                QuoteType::TokenBExactOut,
            ] {
                // Discard the result — pass only requires termination.
                let _ = amm_liquidity(lp, 0, 0, qt, ra, rb);
            }
        }
    }

    #[test]
    fn test_below_min_range_rejected() {
        // Range width below MIN_SQRT_PRICE_RANGE should be rejected
        let result = Position::new(
            1 << 64,
            (1 << 64) + MIN_SQRT_PRICE_RANGE - 1,
            PER_M_DENOMINATOR as u32,
        );
        assert!(result.is_err());
    }

    // Regression for [P2-L-02] — bins that only partially overlap a position's
    // sqrt range must still receive a proportional share of that position's
    // reserves. We verify the property indirectly: the total liquidity
    // distributed across the 32 bins must equal the available reserves on
    // each side. If any partial-overlap bin were dropped (the previous
    // contained-only check), this sum would fall short of `reserves`.
    #[test]
    fn test_partial_overlap_bins_preserve_total_liquidity() {
        let lp = LiquidityType::ConcentratedLiquidity {
            lower_sqrt_price: ((1u128 << 64) as f64 * 0.997f64.sqrt()) as u128,
            upper_sqrt_price: ((1u128 << 64) as f64 * 1.003f64.sqrt()) as u128,
        };
        let reserves_a: u64 = 1_000_000_000_000;
        let reserves_b: u64 = 1_000_000_000_000;

        let ask =
            amm_liquidity(lp, 10, 10, QuoteType::TokenBExactIn, reserves_a, reserves_b).unwrap();
        let bid =
            amm_liquidity(lp, 10, 10, QuoteType::TokenAExactIn, reserves_a, reserves_b).unwrap();

        let total_ask: u64 = ask.as_slice().iter().map(|(_, a)| *a).sum();
        let total_bid: u64 = bid.as_slice().iter().map(|(_, a)| *a).sum();
        assert_eq!(total_ask, reserves_a, "ask side must distribute all of A");
        assert_eq!(total_bid, reserves_b, "bid side must distribute all of B");
    }

    /// Reproduces the partial-overlap coverage gap.
    ///
    /// `bin_position_overlap` is the source of truth for "does this bin
    /// receive any of this position's reserves?" — it returns positive width
    /// whenever the bin and the position have ANY overlap (including the
    /// boundary-straddling partial-overlap case).
    ///
    /// The fix replaces the coverage bitmap predicate with the same
    /// semantics. Every bin that `bin_position_overlap` reports a positive
    /// width for must now be marked as covered. The OLD full-containment
    /// predicate would return `false` here for those partial-overlap bins,
    /// causing under-allocation downstream.
    #[test]
    fn coverage_predicate_includes_partial_overlap_bins() {
        // Bin grid: 32 bins spanning [mid, end] with bin width = MIN_SQRT_PRICE_RANGE.
        // Position upper sits HALF a bin past bin index 5, so bin 5 straddles
        // the boundary (partial overlap), bins 0..=4 are fully contained, and
        // bins 6..=31 are entirely outside the position. Position lower sits
        // below mid so the lower edge does not constrain bin coverage.
        let mid: u128 = 1u128 << 64;
        let bin_width: u128 = MIN_SQRT_PRICE_RANGE;
        let half_bin: u128 = bin_width / 2;
        let position = Position::new(
            mid - bin_width,
            mid + 5 * bin_width + half_bin,
            PER_M_DENOMINATOR as u32,
        )
        .unwrap();
        let end = mid + bin_width * (LIQUIDITY_LEVELS as u128);

        let mut full_containment_count = 0;
        let mut partial_overlap_count = 0;
        let mut no_overlap_count = 0;

        for i in 0..LIQUIDITY_LEVELS {
            let (start_s, end_s) = bin_boundaries(i, mid, end).unwrap();
            let (lo, hi) = if start_s <= end_s {
                (start_s, end_s)
            } else {
                (end_s, start_s)
            };

            // OLD predicate (buggy) — full containment.
            let old_predicate = position.lower_sqrt_price <= lo && hi <= position.upper_sqrt_price;
            // NEW predicate — positive overlap.
            let overlap_start = lo.max(position.lower_sqrt_price);
            let overlap_end = hi.min(position.upper_sqrt_price);
            let new_predicate = overlap_end > overlap_start;
            // bin_position_overlap is the source of truth.
            let overlap_width = bin_position_overlap(i, mid, end, &position).unwrap();

            // Invariant 1: new predicate matches `bin_position_overlap > 0`.
            assert_eq!(
                new_predicate,
                overlap_width > 0,
                "bin {i}: new predicate must agree with bin_position_overlap"
            );

            // Invariant 2: full containment is a STRICT subset of positive overlap.
            if old_predicate {
                assert!(
                    new_predicate,
                    "bin {i}: full-containment implies positive overlap"
                );
            }

            if old_predicate {
                full_containment_count += 1;
            } else if new_predicate {
                partial_overlap_count += 1;
            } else {
                no_overlap_count += 1;
            }
        }

        // Sanity: the constructed scenario actually exercises all three regimes.
        assert!(
            full_containment_count > 0,
            "scenario must include fully-contained bins"
        );
        assert!(
            no_overlap_count > 0,
            "scenario must include no-overlap bins"
        );
        assert!(
            partial_overlap_count > 0,
            "scenario must include at least one partial-overlap bin — \
             this is the case the old predicate dropped"
        );
    }
}