feat(amm): apply trading fees to LP accounting

amm/src/add.rs

@@ -80,17 +80,16 @@ pub fn add_liquidity(
         "Vaults' balances must be at least the reserve amounts"
     );
 
-    // Calculate actual_amounts
-    let ideal_a: u128 = pool_def_data
-        .reserve_a
+    // Quote deposits against live vault balances so newly added LPs do not
+    // receive a share of previously accrued fee surplus.
+    let ideal_a = vault_a_balance
         .checked_mul(max_amount_to_add_token_b)
-        .expect("reserve_a * max_amount_b overflows u128")
-        / pool_def_data.reserve_b;
-    let ideal_b: u128 = pool_def_data
-        .reserve_b
+        .expect("vault_a_balance * max_amount_to_add_token_b overflows u128")
+        / vault_b_balance;
+    let ideal_b = vault_b_balance
         .checked_mul(max_amount_to_add_token_a)
-        .expect("reserve_b * max_amount_a overflows u128")
-        / pool_def_data.reserve_a;
+        .expect("vault_b_balance * max_amount_to_add_token_a overflows u128")
+        / vault_a_balance;
 
     let actual_amount_a = if ideal_a > max_amount_to_add_token_a {
         max_amount_to_add_token_a
@@ -122,12 +121,12 @@ pub fn add_liquidity(
             .liquidity_pool_supply
             .checked_mul(actual_amount_a)
             .expect("liquidity_pool_supply * actual_amount_a overflows u128")
-            / pool_def_data.reserve_a,
+            / vault_a_balance,
         pool_def_data
             .liquidity_pool_supply
             .checked_mul(actual_amount_b)
             .expect("liquidity_pool_supply * actual_amount_b overflows u128")
-            / pool_def_data.reserve_b,
+            / vault_b_balance,
     );
 
     assert!(delta_lp != 0, "Payable LP must be nonzero");

amm/src/remove.rs

@@ -60,7 +60,41 @@ pub fn remove_liquidity(
         "Minimum withdraw amount must be nonzero"
     );
 
-    // 2. Compute withdrawal amounts
+    // 2. Read live vault balances and compute withdrawal amounts
+    let vault_a_token_holding = token_core::TokenHolding::try_from(&vault_a.account.data)
+        .expect("Remove liquidity: AMM Program expects a valid Token Holding Account for Vault A");
+    let token_core::TokenHolding::Fungible {
+        definition_id: _,
+        balance: vault_a_balance,
+    } = vault_a_token_holding
+    else {
+        panic!(
+            "Remove liquidity: AMM Program expects a valid Fungible Token Holding Account for Vault A"
+        );
+    };
+
+    assert!(
+        vault_a_balance >= pool_def_data.reserve_a,
+        "Reserve for Token A exceeds vault balance"
+    );
+
+    let vault_b_token_holding = token_core::TokenHolding::try_from(&vault_b.account.data)
+        .expect("Remove liquidity: AMM Program expects a valid Token Holding Account for Vault B");
+    let token_core::TokenHolding::Fungible {
+        definition_id: _,
+        balance: vault_b_balance,
+    } = vault_b_token_holding
+    else {
+        panic!(
+            "Remove liquidity: AMM Program expects a valid Fungible Token Holding Account for Vault B"
+        );
+    };
+
+    assert!(
+        vault_b_balance >= pool_def_data.reserve_b,
+        "Reserve for Token B exceeds vault balance"
+    );
+
     let user_holding_lp_data = token_core::TokenHolding::try_from(&user_holding_lp.account.data)
         .expect("Remove liquidity: AMM Program expects a valid Token Account for liquidity token");
     let token_core::TokenHolding::Fungible {
@@ -96,46 +130,58 @@ pub fn remove_liquidity(
         "Cannot remove locked minimum liquidity"
     );
 
-    let withdraw_amount_a = pool_def_data
+    // Reserve accounting stays anchored to tracked reserves, while user withdrawals use the
+    // live vault balances so donated surplus is paid out proportionally.
+    let reserve_withdraw_amount_a = pool_def_data
         .reserve_a
         .checked_mul(remove_liquidity_amount)
         .expect("reserve_a * remove_liquidity_amount overflows u128")
         / pool_def_data.liquidity_pool_supply;
-    let withdraw_amount_b = pool_def_data
+    let reserve_withdraw_amount_b = pool_def_data
         .reserve_b
         .checked_mul(remove_liquidity_amount)
         .expect("reserve_b * remove_liquidity_amount overflows u128")
         / pool_def_data.liquidity_pool_supply;
+    let actual_withdraw_amount_a = vault_a_balance
+        .checked_mul(remove_liquidity_amount)
+        .expect("vault_a_balance * remove_liquidity_amount overflows u128")
+        / pool_def_data.liquidity_pool_supply;
+    let actual_withdraw_amount_b = vault_b_balance
+        .checked_mul(remove_liquidity_amount)
+        .expect("vault_b_balance * remove_liquidity_amount overflows u128")
+        / pool_def_data.liquidity_pool_supply;
 
     // 3. Validate and slippage check
     assert!(
-        withdraw_amount_a >= min_amount_to_remove_token_a,
+        actual_withdraw_amount_a >= min_amount_to_remove_token_a,
         "Insufficient minimal withdraw amount (Token A) provided for liquidity amount"
     );
     assert!(
-        withdraw_amount_b >= min_amount_to_remove_token_b,
+        actual_withdraw_amount_b >= min_amount_to_remove_token_b,
         "Insufficient minimal withdraw amount (Token B) provided for liquidity amount"
     );
 
-    // 4. Calculate LP to reduce cap by
-    let delta_lp: u128 = remove_liquidity_amount;
+    // 4. Burn exactly the requested LP amount.
+    let burn_amount_lp = remove_liquidity_amount;
+    let remaining_liquidity = pool_def_data
+        .liquidity_pool_supply
+        .checked_sub(burn_amount_lp)
+        .expect("liquidity_pool_supply - burn_amount_lp underflows");
+    let active = remaining_liquidity != 0;
 
     // 5. Update pool account
     let mut pool_post = pool.account.clone();
     let pool_post_definition = PoolDefinition {
-        liquidity_pool_supply: pool_def_data
-            .liquidity_pool_supply
-            .checked_sub(delta_lp)
-            .expect("liquidity_pool_supply - delta_lp underflows"),
+        liquidity_pool_supply: remaining_liquidity,
         reserve_a: pool_def_data
             .reserve_a
-            .checked_sub(withdraw_amount_a)
-            .expect("reserve_a - withdraw_amount_a underflows"),
+            .checked_sub(reserve_withdraw_amount_a)
+            .expect("reserve_a - reserve_withdraw_amount_a underflows"),
         reserve_b: pool_def_data
             .reserve_b
-            .checked_sub(withdraw_amount_b)
-            .expect("reserve_b - withdraw_amount_b underflows"),
-        active: true,
+            .checked_sub(reserve_withdraw_amount_b)
+            .expect("reserve_b - reserve_withdraw_amount_b underflows"),
+        active,
         ..pool_def_data.clone()
     };
 
@@ -148,7 +194,7 @@ pub fn remove_liquidity(
         token_program_id,
         vec![running_vault_a, user_holding_a.clone()],
         &token_core::Instruction::Transfer {
-            amount_to_transfer: withdraw_amount_a,
+            amount_to_transfer: actual_withdraw_amount_a,
         },
     )
     .with_pda_seeds(vec![compute_vault_pda_seed(
@@ -160,7 +206,7 @@ pub fn remove_liquidity(
         token_program_id,
         vec![running_vault_b, user_holding_b.clone()],
         &token_core::Instruction::Transfer {
-            amount_to_transfer: withdraw_amount_b,
+            amount_to_transfer: actual_withdraw_amount_b,
         },
     )
     .with_pda_seeds(vec![compute_vault_pda_seed(
@@ -174,7 +220,7 @@ pub fn remove_liquidity(
         token_program_id,
         vec![pool_definition_lp_auth, user_holding_lp.clone()],
         &token_core::Instruction::Burn {
-            amount_to_burn: delta_lp,
+            amount_to_burn: burn_amount_lp,
         },
     )
     .with_pda_seeds(vec![compute_liquidity_token_pda_seed(pool.account_id)]);

amm/src/swap.rs

@@ -1,4 +1,4 @@
-use amm_core::assert_supported_fee_tier;
+use amm_core::{assert_supported_fee_tier, FEE_BPS_DENOMINATOR};
 pub use amm_core::{compute_liquidity_token_pda_seed, compute_vault_pda_seed, PoolDefinition};
 use nssa_core::{
     account::{AccountId, AccountWithMetadata, Data},
@@ -127,6 +127,7 @@ pub fn swap_exact_input(
                 user_holding_b.clone(),
                 swap_amount_in,
                 min_amount_out,
+                pool_def_data.fees,
                 pool_def_data.reserve_a,
                 pool_def_data.reserve_b,
                 pool.account_id,
@@ -141,6 +142,7 @@ pub fn swap_exact_input(
                 user_holding_a.clone(),
                 swap_amount_in,
                 min_amount_out,
+                pool_def_data.fees,
                 pool_def_data.reserve_b,
                 pool_def_data.reserve_a,
                 pool.account_id,
@@ -175,19 +177,31 @@ fn swap_logic(
     user_withdraw: AccountWithMetadata,
     swap_amount_in: u128,
     min_amount_out: u128,
+    fee_bps: u128,
     reserve_deposit_vault_amount: u128,
     reserve_withdraw_vault_amount: u128,
     pool_id: AccountId,
 ) -> (Vec<ChainedCall>, u128, u128) {
-    // Compute withdraw amount
-    // Maintains pool constant product
-    // k = pool_def_data.reserve_a * pool_def_data.reserve_b;
+    let fee_multiplier = FEE_BPS_DENOMINATOR
+        .checked_sub(fee_bps)
+        .expect("fee_bps exceeds fee denominator");
+    let effective_amount_in = swap_amount_in
+        .checked_mul(fee_multiplier)
+        .expect("swap_amount_in * fee_multiplier overflows u128")
+        / FEE_BPS_DENOMINATOR;
+    assert!(
+        effective_amount_in != 0,
+        "Effective swap amount should be nonzero"
+    );
+
+    // Compute withdraw amount from fee-adjusted reserves while leaving the fee
+    // portion behind as vault surplus for LPs.
     let withdraw_amount = reserve_withdraw_vault_amount
-        .checked_mul(swap_amount_in)
-        .expect("reserve * amount_in overflows u128")
+        .checked_mul(effective_amount_in)
+        .expect("reserve_withdraw_vault_amount * effective_amount_in overflows u128")
         / reserve_deposit_vault_amount
-            .checked_add(swap_amount_in)
-            .expect("reserve + swap_amount_in overflows u128");
+            .checked_add(effective_amount_in)
+            .expect("reserve_deposit_vault_amount + effective_amount_in overflows u128");
 
     // Slippage check
     assert!(
@@ -228,7 +242,7 @@ fn swap_logic(
         .with_pda_seeds(vec![pda_seed]),
     );
 
-    (chained_calls, swap_amount_in, withdraw_amount)
+    (chained_calls, effective_amount_in, withdraw_amount)
 }
 
 #[expect(clippy::too_many_arguments, reason = "TODO: Fix later")]
@@ -254,6 +268,7 @@ pub fn swap_exact_output(
                 user_holding_b.clone(),
                 exact_amount_out,
                 max_amount_in,
+                pool_def_data.fees,
                 pool_def_data.reserve_a,
                 pool_def_data.reserve_b,
                 pool.account_id,
@@ -268,6 +283,7 @@ pub fn swap_exact_output(
                 user_holding_a.clone(),
                 exact_amount_out,
                 max_amount_in,
+                pool_def_data.fees,
                 pool_def_data.reserve_b,
                 pool_def_data.reserve_a,
                 pool.account_id,
@@ -302,6 +318,7 @@ fn exact_output_swap_logic(
     user_withdraw: AccountWithMetadata,
     exact_amount_out: u128,
     max_amount_in: u128,
+    fee_bps: u128,
     reserve_deposit_vault_amount: u128,
     reserve_withdraw_vault_amount: u128,
     pool_id: AccountId,
@@ -317,10 +334,21 @@ fn exact_output_swap_logic(
 
     // Compute deposit amount using ceiling division
     // Formula: amount_in = ceil(reserve_in * exact_amount_out / (reserve_out - exact_amount_out))
-    let deposit_amount = reserve_deposit_vault_amount
+    let effective_deposit_amount = reserve_deposit_vault_amount
         .checked_mul(exact_amount_out)
         .expect("reserve * amount_out overflows u128")
-        .div_ceil(reserve_withdraw_vault_amount - exact_amount_out);
+        .div_ceil(
+            reserve_withdraw_vault_amount
+                .checked_sub(exact_amount_out)
+                .expect("reserve_withdraw_vault_amount - exact_amount_out underflows"),
+        );
+    let fee_multiplier = FEE_BPS_DENOMINATOR
+        .checked_sub(fee_bps)
+        .expect("fee_bps exceeds fee denominator");
+    let deposit_amount = effective_deposit_amount
+        .checked_mul(FEE_BPS_DENOMINATOR)
+        .expect("effective_deposit_amount * fee denominator overflows u128")
+        .div_ceil(fee_multiplier);
 
     // Slippage check
     assert!(
@@ -360,5 +388,5 @@ fn exact_output_swap_logic(
         .with_pda_seeds(vec![pda_seed]),
     );
 
-    (chained_calls, deposit_amount, exact_amount_out)
+    (chained_calls, effective_deposit_amount, exact_amount_out)
 }