diff --git a/programs/amm/core/src/lib.rs b/programs/amm/core/src/lib.rs index 8859a6a..3351493 100644 --- a/programs/amm/core/src/lib.rs +++ b/programs/amm/core/src/lib.rs @@ -334,6 +334,37 @@ pub fn mul_div_ceil(a: u128, b: u128, c: u128) -> u128 { u128::try_from(result).expect("mul_div_ceil result exceeds u128") } +/// The constant-product output for a `SwapExactInput`, matching the AMM's on-chain +/// pricing exactly — used by both `amm_program::swap` and the off-chain swap quote, +/// so the preview and the executed trade agree. Fee-adjusts the input, then applies +/// `reserve_out * effective / (reserve_in + effective)`. +/// +/// Returns `(effective_amount_in, amount_out)`; both are `0` when the fee-adjusted +/// input rounds to zero. Saturating: an out-of-range fee, or the impossible +/// `reserve_in + effective` overflow (would need a reserve near `u128::MAX`), +/// degrades to `0` output rather than panicking. Callers validate the fee tier and +/// enforce their own nonzero / min-out checks. +#[must_use] +pub fn swap_exact_in_amounts( + amount_in: u128, + reserve_in: u128, + reserve_out: u128, + fee_bps: u128, +) -> (u128, u128) { + let fee_multiplier = FEE_BPS_DENOMINATOR.saturating_sub(fee_bps); + let effective_amount_in = mul_div_floor(amount_in, fee_multiplier, FEE_BPS_DENOMINATOR); + let reserve_plus_effective = match reserve_in.checked_add(effective_amount_in) { + Some(v) => v, + None => return (effective_amount_in, 0), + }; + let amount_out = if reserve_plus_effective == 0 { + 0 + } else { + mul_div_floor(reserve_out, effective_amount_in, reserve_plus_effective) + }; + (effective_amount_in, amount_out) +} + /// `floor(sqrt(a * b))` computed in U256 so the `a * b` product can't overflow u128. /// /// # Panics @@ -633,6 +664,21 @@ mod tests { let _ = mul_div_floor(1, 2, 0); } + #[test] + fn swap_exact_in_amounts_matches_constant_product() { + // 0.30% fee, reserves 1_000_000 in / 2_000_000 out, amount_in 10_000. + let (eff, out) = swap_exact_in_amounts(10_000, 1_000_000, 2_000_000, 30); + let expected_eff = 10_000 * (10_000 - 30) / 10_000; + let expected_out = 2_000_000 * expected_eff / (1_000_000 + expected_eff); + assert_eq!((eff, out), (expected_eff, expected_out)); + + // A tiny input can fee-round the effective input to zero → zero output. + assert_eq!(swap_exact_in_amounts(1, 1_000_000, 2_000_000, 30), (0, 0)); + + // Degenerate reserves saturate rather than dividing by zero. + assert_eq!(swap_exact_in_amounts(1, 0, 0, 30), (0, 0)); + } + #[test] fn mul_div_ceil_small_cases() { assert_eq!(mul_div_ceil(6, 7, 3), 14); diff --git a/programs/amm/src/swap.rs b/programs/amm/src/swap.rs index 1f10252..37bcc64 100644 --- a/programs/amm/src/swap.rs +++ b/programs/amm/src/swap.rs @@ -1,7 +1,7 @@ use amm_core::{ assert_supported_fee_tier, compute_config_pda, compute_pool_pda_seed, mul_div_ceil, - mul_div_floor, read_vault_fungible_balances, spot_price_q64_64, AmmConfig, FEE_BPS_DENOMINATOR, - MINIMUM_LIQUIDITY, + read_vault_fungible_balances, spot_price_q64_64, swap_exact_in_amounts, AmmConfig, + FEE_BPS_DENOMINATOR, MINIMUM_LIQUIDITY, }; pub use amm_core::{compute_liquidity_token_pda_seed, compute_vault_pda_seed, PoolDefinition}; use clock_core::CLOCK_01_PROGRAM_ACCOUNT_ID; @@ -291,29 +291,21 @@ fn swap_logic( reserve_withdraw_vault_amount: u128, pool_id: AccountId, ) -> (Vec, u128, u128) { - let fee_multiplier = FEE_BPS_DENOMINATOR - .checked_sub(fee_bps) - .expect("fee_bps exceeds fee denominator"); - // floor(swap_amount_in * fee_multiplier / FEE_BPS_DENOMINATOR), product widened to U256. - let effective_amount_in = mul_div_floor(swap_amount_in, fee_multiplier, FEE_BPS_DENOMINATOR); + // Fee-adjust the input and price via constant product. Shared with the + // off-chain swap quote (`amm_core::swap_exact_in_amounts`) so the preview and + // the executed trade agree exactly. The recorded pool reserves are updated + // later with the full `swap_amount_in`, so LP fees accrue inside `reserve_*` + // via invariant growth rather than as a vault-balance surplus over `reserve_*`. + let (effective_amount_in, withdraw_amount) = swap_exact_in_amounts( + swap_amount_in, + reserve_deposit_vault_amount, + reserve_withdraw_vault_amount, + fee_bps, + ); assert!( effective_amount_in != 0, "Effective swap amount should be nonzero" ); - // Compute the withdraw amount using the fee-adjusted input for pricing. - // The recorded pool reserves are updated later with the full - // `swap_amount_in`, so LP fees accrue inside `reserve_*` via invariant - // growth rather than as a separate vault balance surplus over `reserve_*`. - // The denominator sum stays u128 (overflows only near u128::MAX, an unstorable reserve); - // only the `reserve * effective` product is widened to U256. - let reserve_plus_effective = reserve_deposit_vault_amount - .checked_add(effective_amount_in) - .expect("reserve + effective_amount_in overflows u128"); - let withdraw_amount = mul_div_floor( - reserve_withdraw_vault_amount, - effective_amount_in, - reserve_plus_effective, - ); // Slippage check assert!(