From 8fba03d106dfa4b6d2f95252a45cceed556a9c0a Mon Sep 17 00:00:00 2001
From: r4bbit <445106+0x-r4bbit@users.noreply.github.com>
Date: Tue, 4 Aug 2026 19:04:02 +0200
Subject: [PATCH] refactor(amm): extract swap-input formula into amm_core

Lift the inverse constant-product SwapExactOutput math out of the guest's
exact_output_swap_logic into amm_core::swap_exact_out_amounts(amount_out,
reserve_in, reserve_out, fee_bps) -> Option<(effective_in, required_in)>:
ceil(reserve_in * amount_out / (reserve_out - amount_out)) lifted through the
fee via mul_div_ceil. exact_output_swap_logic now calls it and keeps its
nonzero/exceeds-reserve asserts. Behavior-preserving (the panic-message tests
still pass); None (out >= reserve or zero fee multiplier) surfaces as the
existing expect.

Makes the on-chain exact-output pricing one reusable function so the
off-chain quote can produce byte-identical required-input figures instead of
re-deriving the inverse formula. Mirrors swap_exact_in_amounts.
---
 programs/amm/core/src/lib.rs | 63 ++++++++++++++++++++++++++++++++++++
 programs/amm/src/swap.rs     | 37 +++++++--------------
 2 files changed, 75 insertions(+), 25 deletions(-)

diff --git a/programs/amm/core/src/lib.rs b/programs/amm/core/src/lib.rs
index ac09a86..6951535 100644
--- a/programs/amm/core/src/lib.rs
+++ b/programs/amm/core/src/lib.rs
@@ -409,6 +409,42 @@ pub fn swap_exact_in_amounts(
     (effective_amount_in, amount_out)
 }
 
+/// The input amounts for a `SwapExactOutput`: the fee-adjusted (effective) input
+/// and the gross input required to receive exactly `amount_out`, matching the AMM's
+/// on-chain pricing — used by both `amm_program::swap` and the off-chain
+/// exact-output quote. Solves the constant product for the input, then lifts it
+/// back through the fee (both steps round up, so the pool never comes up short):
+/// `required_in = ceil(ceil(reserve_in * amount_out / (reserve_out - amount_out))
+/// * FEE_DENOM / (FEE_DENOM - fee_bps))`.
+///
+/// Returns `None` when the trade is unfulfillable: `amount_out >= reserve_out` (you
+/// can't withdraw the whole pool), `reserve_in == 0` (an empty input reserve has no
+/// solution for a positive output — without this it would round to a free
+/// `(0, 0)`), or a degenerate `fee_bps >= FEE_DENOM`. Callers enforce their own
+/// nonzero / max-in checks.
+#[must_use]
+pub fn swap_exact_out_amounts(
+    amount_out: u128,
+    reserve_in: u128,
+    reserve_out: u128,
+    fee_bps: u128,
+) -> Option<(u128, u128)> {
+    let denominator = reserve_out.checked_sub(amount_out)?;
+    if denominator == 0 {
+        return None; // amount_out == reserve_out: draining the pool is impossible
+    }
+    if reserve_in == 0 {
+        return None; // no input reserve: the inverse curve has no solution for a positive output
+    }
+    let effective_in_min = mul_div_ceil(reserve_in, amount_out, denominator);
+    let fee_multiplier = FEE_BPS_DENOMINATOR.checked_sub(fee_bps)?;
+    if fee_multiplier == 0 {
+        return None; // fee_bps == FEE_DENOM (impossible for a supported tier)
+    }
+    let required_in = mul_div_ceil(effective_in_min, FEE_BPS_DENOMINATOR, fee_multiplier);
+    Some((effective_in_min, required_in))
+}
+
 /// `floor(sqrt(a * b))` computed in U256 so the `a * b` product can't overflow u128.
 ///
 /// # Panics
@@ -743,6 +779,33 @@ mod tests {
         assert_eq!(price_impact_bps(1, 0, 1_000_000, 0), 0);
     }
 
+    #[test]
+    fn swap_exact_out_amounts_inverts_the_constant_product() {
+        // 0.30% fee, reserves 1_000_000 in / 2_000_000 out, want 10_000 out.
+        let (eff_min, required_in) =
+            swap_exact_out_amounts(10_000, 1_000_000, 2_000_000, 30).unwrap();
+        // effective_in >= ceil(reserve_in * out / (reserve_out - out)); then lift through fee.
+        let expected_eff = (1_000_000u128 * 10_000).div_ceil(2_000_000 - 10_000);
+        let expected_in = (expected_eff * 10_000).div_ceil(10_000 - 30);
+        assert_eq!((eff_min, required_in), (expected_eff, expected_in));
+
+        // Round-trips with the forward quote: paying required_in yields at least the ask.
+        let (_, out) = swap_exact_in_amounts(required_in, 1_000_000, 2_000_000, 30);
+        assert!(out >= 10_000);
+
+        // Unfulfillable: can't withdraw the whole pool (or more).
+        assert_eq!(
+            swap_exact_out_amounts(2_000_000, 1_000_000, 2_000_000, 30),
+            None
+        );
+        assert_eq!(
+            swap_exact_out_amounts(2_000_001, 1_000_000, 2_000_000, 30),
+            None
+        );
+        // No input reserve → no valid input for a positive output (not a free (0, 0)).
+        assert_eq!(swap_exact_out_amounts(10_000, 0, 2_000_000, 30), None);
+    }
+
     #[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 37bcc64..794198d 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,
-    read_vault_fungible_balances, spot_price_q64_64, swap_exact_in_amounts, AmmConfig,
-    FEE_BPS_DENOMINATOR, MINIMUM_LIQUIDITY,
+    assert_supported_fee_tier, compute_config_pda, compute_pool_pda_seed,
+    read_vault_fungible_balances, spot_price_q64_64, swap_exact_in_amounts, swap_exact_out_amounts,
+    AmmConfig, MINIMUM_LIQUIDITY,
 };
 pub use amm_core::{compute_liquidity_token_pda_seed, compute_vault_pda_seed, PoolDefinition};
 use clock_core::CLOCK_01_PROGRAM_ACCOUNT_ID;
@@ -514,29 +514,16 @@ fn exact_output_swap_logic(
         "Exact amount out exceeds reserve"
     );
 
-    // Compute the minimum effective input required to achieve exact_amount_out
-    // using the same floor-rounded fee application as swap_exact_input.
-    //
-    // Solve constant product for effective_in (fee already removed):
-    //   effective_in >= ceil(reserve_in * amount_out / (reserve_out - amount_out))
-    // ceil(reserve_in * amount_out / (reserve_out - amount_out)). The `reserve_in * amount_out`
-    // product is widened to U256; the denominator is a subtraction that stays u128.
-    let effective_in_denominator = reserve_withdraw_vault_amount
-        .checked_sub(exact_amount_out)
-        .expect("reserve_out - amount_out underflows");
-    let effective_in_min = mul_div_ceil(
-        reserve_deposit_vault_amount,
+    // Required gross input via the shared amm_core::swap_exact_out_amounts (same
+    // pricing as the off-chain exact-output quote). The `amount_out < reserve`
+    // guard above means it always resolves.
+    let (_, deposit_amount) = swap_exact_out_amounts(
         exact_amount_out,
-        effective_in_denominator,
-    );
-
-    // Lift back to gross input so that
-    //   floor(gross_in * (FEE_DENOM - fee) / FEE_DENOM) >= effective_in_min
-    // ceil(effective_in_min * FEE_BPS_DENOMINATOR / fee_multiplier), product widened to U256.
-    let fee_multiplier = FEE_BPS_DENOMINATOR
-        .checked_sub(fee_bps)
-        .expect("fee_bps exceeds fee denominator");
-    let deposit_amount = mul_div_ceil(effective_in_min, FEE_BPS_DENOMINATOR, fee_multiplier);
+        reserve_deposit_vault_amount,
+        reserve_withdraw_vault_amount,
+        fee_bps,
+    )
+    .expect("swap exact output: reserves and fee must yield a valid input");
 
     // Slippage check
     assert!(