diff --git a/amm/src/add.rs b/amm/src/add.rs index 91f5962..a2b8881 100644 --- a/amm/src/add.rs +++ b/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"); diff --git a/amm/src/remove.rs b/amm/src/remove.rs index ab1f7d8..a5a1042 100644 --- a/amm/src/remove.rs +++ b/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)]); diff --git a/amm/src/swap.rs b/amm/src/swap.rs index fe5c544..27cd41e 100644 --- a/amm/src/swap.rs +++ b/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, 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) } diff --git a/amm/src/tests.rs b/amm/src/tests.rs index 70b8434..49304f1 100644 --- a/amm/src/tests.rs +++ b/amm/src/tests.rs @@ -4,8 +4,9 @@ use std::num::NonZero; use amm_core::{ compute_liquidity_token_pda, compute_liquidity_token_pda_seed, compute_lp_lock_holding_pda, - compute_pool_pda, compute_vault_pda, compute_vault_pda_seed, PoolDefinition, FEE_TIER_BPS_1, - FEE_TIER_BPS_100, FEE_TIER_BPS_30, FEE_TIER_BPS_5, MINIMUM_LIQUIDITY, + compute_pool_pda, compute_vault_pda, compute_vault_pda_seed, PoolDefinition, + FEE_BPS_DENOMINATOR, FEE_TIER_BPS_1, FEE_TIER_BPS_100, FEE_TIER_BPS_30, FEE_TIER_BPS_5, + MINIMUM_LIQUIDITY, }; use nssa_core::{ account::{Account, AccountId, AccountWithMetadata, Data, Nonce}, @@ -124,8 +125,18 @@ impl BalanceForTests { BalanceForTests::lp_supply_init() - MINIMUM_LIQUIDITY } + fn effective_swap_amount_in_a() -> u128 { + BalanceForTests::add_max_amount_a() * (FEE_BPS_DENOMINATOR - BalanceForTests::fee_tier()) + / FEE_BPS_DENOMINATOR + } + + fn effective_swap_amount_in_b() -> u128 { + BalanceForTests::add_max_amount_b() * (FEE_BPS_DENOMINATOR - BalanceForTests::fee_tier()) + / FEE_BPS_DENOMINATOR + } + fn vault_a_swap_test_1() -> u128 { - BalanceForTests::vault_a_reserve_init() + BalanceForTests::add_max_amount_a() + BalanceForTests::vault_a_reserve_init() + BalanceForTests::effective_swap_amount_in_a() } fn vault_a_swap_test_2() -> u128 { @@ -137,7 +148,7 @@ impl BalanceForTests { } fn vault_b_swap_test_2() -> u128 { - BalanceForTests::vault_b_reserve_init() + BalanceForTests::add_max_amount_b() + BalanceForTests::vault_b_reserve_init() + BalanceForTests::effective_swap_amount_in_b() } fn min_amount_out() -> u128 { @@ -169,6 +180,76 @@ impl BalanceForTests { 166 } + fn vault_a_balance_with_surplus() -> u128 { + BalanceForTests::vault_a_reserve_init() + (BalanceForTests::vault_a_reserve_init() / 10) + } + + fn vault_b_balance_with_surplus() -> u128 { + BalanceForTests::vault_b_reserve_init() + (BalanceForTests::vault_b_reserve_init() / 10) + } + + fn exact_output_effective_amount_in_token_a() -> u128 { + 1_000_u128 + .checked_mul(BalanceForTests::max_amount_in()) + .expect("exact output reserve_a * max_amount_in overflows u128") + .div_ceil(500 - BalanceForTests::max_amount_in()) + } + + fn exact_output_effective_amount_in_token_b() -> u128 { + 500_u128 + .checked_mul(285) + .expect("exact output reserve_b * exact_amount_out overflows u128") + .div_ceil(1_000 - 285) + } + + fn exact_output_deposit_amount_token_a() -> u128 { + BalanceForTests::exact_output_effective_amount_in_token_a() + .checked_mul(FEE_BPS_DENOMINATOR) + .expect("effective amount in * fee denominator overflows u128") + .div_ceil(FEE_BPS_DENOMINATOR - BalanceForTests::fee_tier()) + } + + fn exact_output_deposit_amount_token_b() -> u128 { + BalanceForTests::exact_output_effective_amount_in_token_b() + .checked_mul(FEE_BPS_DENOMINATOR) + .expect("effective amount in * fee denominator overflows u128") + .div_ceil(FEE_BPS_DENOMINATOR - BalanceForTests::fee_tier()) + } + + fn reserve_a_add_with_surplus() -> u128 { + BalanceForTests::vault_a_reserve_init() + + BalanceForTests::add_successful_amount_a_with_surplus() + } + + fn reserve_b_add_with_surplus() -> u128 { + BalanceForTests::vault_b_reserve_init() + + BalanceForTests::add_successful_amount_b_with_surplus() + } + + fn add_successful_amount_a_with_surplus() -> u128 { + (BalanceForTests::vault_a_balance_with_surplus() * BalanceForTests::add_max_amount_b()) + / BalanceForTests::vault_b_balance_with_surplus() + } + + fn add_successful_amount_b_with_surplus() -> u128 { + BalanceForTests::add_max_amount_b() + } + + fn lp_supply_with_surplus() -> u128 { + BalanceForTests::lp_supply_init() + BalanceForTests::lp_mint_with_surplus() + } + + fn lp_mint_with_surplus() -> u128 { + std::cmp::min( + BalanceForTests::lp_supply_init() + * BalanceForTests::add_successful_amount_a_with_surplus() + / BalanceForTests::vault_a_balance_with_surplus(), + BalanceForTests::lp_supply_init() + * BalanceForTests::add_successful_amount_b_with_surplus() + / BalanceForTests::vault_b_balance_with_surplus(), + ) + } + fn vault_a_remove_successful() -> u128 { BalanceForTests::vault_a_reserve_init() - BalanceForTests::remove_actual_a_successful() } @@ -178,13 +259,15 @@ impl BalanceForTests { } fn swap_amount_out_b() -> u128 { - (BalanceForTests::vault_b_reserve_init() * BalanceForTests::add_max_amount_a()) - / (BalanceForTests::vault_a_reserve_init() + BalanceForTests::add_max_amount_a()) + (BalanceForTests::vault_b_reserve_init() * BalanceForTests::effective_swap_amount_in_a()) + / (BalanceForTests::vault_a_reserve_init() + + BalanceForTests::effective_swap_amount_in_a()) } fn swap_amount_out_a() -> u128 { - (BalanceForTests::vault_a_reserve_init() * BalanceForTests::add_max_amount_b()) - / (BalanceForTests::vault_b_reserve_init() + BalanceForTests::add_max_amount_b()) + (BalanceForTests::vault_a_reserve_init() * BalanceForTests::effective_swap_amount_in_b()) + / (BalanceForTests::vault_b_reserve_init() + + BalanceForTests::effective_swap_amount_in_b()) } fn add_delta_lp_successful() -> u128 { @@ -208,6 +291,20 @@ impl BalanceForTests { fn remove_lp_supply_successful() -> u128 { BalanceForTests::lp_supply_init() - BalanceForTests::remove_amount_lp() } + + fn remove_actual_a_with_surplus() -> u128 { + (BalanceForTests::vault_a_balance_with_surplus() * BalanceForTests::remove_amount_lp()) + / BalanceForTests::lp_supply_init() + } + + fn remove_actual_b_with_surplus() -> u128 { + (BalanceForTests::vault_b_balance_with_surplus() * BalanceForTests::remove_amount_lp()) + / BalanceForTests::lp_supply_init() + } + + fn remove_min_amount_b_surplus() -> u128 { + BalanceForTests::remove_actual_b_with_surplus() + } } impl ChainedCallForTests { @@ -276,7 +373,7 @@ impl ChainedCallForTests { } fn cc_swap_exact_output_token_a_test_1() -> ChainedCall { - let swap_amount: u128 = 498; + let swap_amount = BalanceForTests::exact_output_deposit_amount_token_a(); ChainedCall::new( TOKEN_PROGRAM_ID, @@ -329,7 +426,7 @@ impl ChainedCallForTests { } fn cc_swap_exact_output_token_b_test_2() -> ChainedCall { - let swap_amount: u128 = 200; + let swap_amount = BalanceForTests::exact_output_deposit_amount_token_b(); ChainedCall::new( TOKEN_PROGRAM_ID, @@ -388,6 +485,51 @@ impl ChainedCallForTests { )]) } + fn cc_add_token_a_with_surplus() -> ChainedCall { + ChainedCall::new( + TOKEN_PROGRAM_ID, + vec![ + AccountWithMetadataForTests::user_holding_a(), + AccountWithMetadataForTests::vault_a_with_surplus(), + ], + &token_core::Instruction::Transfer { + amount_to_transfer: BalanceForTests::add_successful_amount_a_with_surplus(), + }, + ) + } + + fn cc_add_token_b_with_surplus() -> ChainedCall { + ChainedCall::new( + TOKEN_PROGRAM_ID, + vec![ + AccountWithMetadataForTests::user_holding_b(), + AccountWithMetadataForTests::vault_b_with_surplus(), + ], + &token_core::Instruction::Transfer { + amount_to_transfer: BalanceForTests::add_successful_amount_b_with_surplus(), + }, + ) + } + + fn cc_add_pool_lp_with_surplus() -> ChainedCall { + let mut pool_lp_auth = AccountWithMetadataForTests::pool_lp_init(); + pool_lp_auth.is_authorized = true; + + ChainedCall::new( + TOKEN_PROGRAM_ID, + vec![ + pool_lp_auth, + AccountWithMetadataForTests::user_holding_lp_init(), + ], + &token_core::Instruction::Mint { + amount_to_mint: BalanceForTests::lp_mint_with_surplus(), + }, + ) + .with_pda_seeds(vec![compute_liquidity_token_pda_seed( + IdForTests::pool_definition_id(), + )]) + } + fn cc_remove_token_a() -> ChainedCall { let mut vault_a_auth = AccountWithMetadataForTests::vault_a_init(); vault_a_auth.is_authorized = true; @@ -441,6 +583,40 @@ impl ChainedCallForTests { )]) } + fn cc_remove_token_a_with_surplus() -> ChainedCall { + let mut vault_a_auth = AccountWithMetadataForTests::vault_a_with_surplus(); + vault_a_auth.is_authorized = true; + + ChainedCall::new( + TOKEN_PROGRAM_ID, + vec![vault_a_auth, AccountWithMetadataForTests::user_holding_a()], + &token_core::Instruction::Transfer { + amount_to_transfer: BalanceForTests::remove_actual_a_with_surplus(), + }, + ) + .with_pda_seeds(vec![compute_vault_pda_seed( + IdForTests::pool_definition_id(), + IdForTests::token_a_definition_id(), + )]) + } + + fn cc_remove_token_b_with_surplus() -> ChainedCall { + let mut vault_b_auth = AccountWithMetadataForTests::vault_b_with_surplus(); + vault_b_auth.is_authorized = true; + + ChainedCall::new( + TOKEN_PROGRAM_ID, + vec![vault_b_auth, AccountWithMetadataForTests::user_holding_b()], + &token_core::Instruction::Transfer { + amount_to_transfer: BalanceForTests::remove_actual_b_with_surplus(), + }, + ) + .with_pda_seeds(vec![compute_vault_pda_seed( + IdForTests::pool_definition_id(), + IdForTests::token_b_definition_id(), + )]) + } + fn cc_new_definition_token_a() -> ChainedCall { ChainedCall::new( TOKEN_PROGRAM_ID, @@ -619,6 +795,38 @@ impl AccountWithMetadataForTests { } } + fn vault_a_with_surplus() -> AccountWithMetadata { + AccountWithMetadata { + account: Account { + program_owner: TOKEN_PROGRAM_ID, + balance: 0u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: IdForTests::token_a_definition_id(), + balance: BalanceForTests::vault_a_balance_with_surplus(), + }), + nonce: Nonce(0), + }, + is_authorized: true, + account_id: IdForTests::vault_a_id(), + } + } + + fn vault_b_with_surplus() -> AccountWithMetadata { + AccountWithMetadata { + account: Account { + program_owner: TOKEN_PROGRAM_ID, + balance: 0u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: IdForTests::token_b_definition_id(), + balance: BalanceForTests::vault_b_balance_with_surplus(), + }), + nonce: Nonce(0), + }, + is_authorized: true, + account_id: IdForTests::vault_b_id(), + } + } + fn vault_a_init_high() -> AccountWithMetadata { AccountWithMetadata { account: Account { @@ -1068,8 +1276,8 @@ impl AccountWithMetadataForTests { vault_b_id: IdForTests::vault_b_id(), liquidity_pool_id: IdForTests::token_lp_definition_id(), liquidity_pool_supply: BalanceForTests::lp_supply_init(), - reserve_a: 1498_u128, - reserve_b: 334_u128, + reserve_a: 1_000 + BalanceForTests::exact_output_effective_amount_in_token_a(), + reserve_b: 500 - BalanceForTests::max_amount_in(), fees: BalanceForTests::fee_tier(), active: true, }), @@ -1092,8 +1300,8 @@ impl AccountWithMetadataForTests { vault_b_id: IdForTests::vault_b_id(), liquidity_pool_id: IdForTests::token_lp_definition_id(), liquidity_pool_supply: BalanceForTests::lp_supply_init(), - reserve_a: 715_u128, - reserve_b: 700_u128, + reserve_a: 1_000 - 285, + reserve_b: 500 + BalanceForTests::exact_output_effective_amount_in_token_b(), fees: BalanceForTests::fee_tier(), active: true, }), @@ -1152,6 +1360,54 @@ impl AccountWithMetadataForTests { } } + fn pool_definition_add_with_surplus_successful() -> AccountWithMetadata { + AccountWithMetadata { + account: Account { + program_owner: ProgramId::default(), + balance: 0u128, + data: Data::from(&PoolDefinition { + definition_token_a_id: IdForTests::token_a_definition_id(), + definition_token_b_id: IdForTests::token_b_definition_id(), + vault_a_id: IdForTests::vault_a_id(), + vault_b_id: IdForTests::vault_b_id(), + liquidity_pool_id: IdForTests::token_lp_definition_id(), + liquidity_pool_supply: BalanceForTests::lp_supply_with_surplus(), + reserve_a: BalanceForTests::reserve_a_add_with_surplus(), + reserve_b: BalanceForTests::reserve_b_add_with_surplus(), + fees: BalanceForTests::fee_tier(), + active: true, + }), + nonce: Nonce(0), + }, + is_authorized: true, + account_id: IdForTests::pool_definition_id(), + } + } + + fn pool_definition_init_low_balances() -> AccountWithMetadata { + AccountWithMetadata { + account: Account { + program_owner: ProgramId::default(), + balance: 0u128, + data: Data::from(&PoolDefinition { + definition_token_a_id: IdForTests::token_a_definition_id(), + definition_token_b_id: IdForTests::token_b_definition_id(), + vault_a_id: IdForTests::vault_a_id(), + vault_b_id: IdForTests::vault_b_id(), + liquidity_pool_id: IdForTests::token_lp_definition_id(), + liquidity_pool_supply: BalanceForTests::vault_a_reserve_low(), + reserve_a: BalanceForTests::vault_a_reserve_low(), + reserve_b: BalanceForTests::vault_b_reserve_low(), + fees: BalanceForTests::fee_tier(), + active: true, + }), + nonce: Nonce(0), + }, + is_authorized: true, + account_id: IdForTests::pool_definition_id(), + } + } + fn pool_definition_remove_successful() -> AccountWithMetadata { AccountWithMetadata { account: Account { @@ -1581,6 +1837,46 @@ fn test_call_add_liquidity_chained_call_successsful() { assert!(chained_call_lp == ChainedCallForTests::cc_add_pool_lp()); } +#[test] +fn test_call_add_liquidity_with_fee_surplus_preserves_existing_lp_value() { + let (post_states, chained_calls) = add_liquidity( + AccountWithMetadataForTests::pool_definition_init(), + AccountWithMetadataForTests::vault_a_with_surplus(), + AccountWithMetadataForTests::vault_b_with_surplus(), + AccountWithMetadataForTests::pool_lp_init(), + AccountWithMetadataForTests::user_holding_a(), + AccountWithMetadataForTests::user_holding_b(), + AccountWithMetadataForTests::user_holding_lp_init(), + NonZero::new(BalanceForTests::add_min_amount_lp()).unwrap(), + BalanceForTests::add_max_amount_a(), + BalanceForTests::add_max_amount_b(), + ); + + let pool_post = post_states[0].clone(); + + assert!( + AccountWithMetadataForTests::pool_definition_add_with_surplus_successful().account + == *pool_post.account() + ); + + let chained_call_lp = chained_calls[0].clone(); + let chained_call_b = chained_calls[1].clone(); + let chained_call_a = chained_calls[2].clone(); + + assert_eq!( + chained_call_a, + ChainedCallForTests::cc_add_token_a_with_surplus() + ); + assert_eq!( + chained_call_b, + ChainedCallForTests::cc_add_token_b_with_surplus() + ); + assert_eq!( + chained_call_lp, + ChainedCallForTests::cc_add_pool_lp_with_surplus() + ); +} + #[should_panic(expected = "Vault A was not provided")] #[test] fn test_call_remove_liquidity_vault_a_omitted() { @@ -1761,6 +2057,40 @@ fn test_call_remove_liquidity_min_bal_zero_2() { ); } +#[should_panic(expected = "Reserve for Token A exceeds vault balance")] +#[test] +fn test_call_remove_liquidity_reserves_vault_mismatch_1() { + let _post_states = remove_liquidity( + AccountWithMetadataForTests::pool_definition_init(), + AccountWithMetadataForTests::vault_a_init_low(), + AccountWithMetadataForTests::vault_b_init(), + AccountWithMetadataForTests::pool_lp_init(), + AccountWithMetadataForTests::user_holding_a(), + AccountWithMetadataForTests::user_holding_b(), + AccountWithMetadataForTests::user_holding_lp_init(), + NonZero::new(BalanceForTests::remove_amount_lp()).unwrap(), + BalanceForTests::remove_min_amount_a(), + BalanceForTests::remove_min_amount_b(), + ); +} + +#[should_panic(expected = "Reserve for Token B exceeds vault balance")] +#[test] +fn test_call_remove_liquidity_reserves_vault_mismatch_2() { + let _post_states = remove_liquidity( + AccountWithMetadataForTests::pool_definition_init(), + AccountWithMetadataForTests::vault_a_init(), + AccountWithMetadataForTests::vault_b_init_low(), + AccountWithMetadataForTests::pool_lp_init(), + AccountWithMetadataForTests::user_holding_a(), + AccountWithMetadataForTests::user_holding_b(), + AccountWithMetadataForTests::user_holding_lp_init(), + NonZero::new(BalanceForTests::remove_amount_lp()).unwrap(), + BalanceForTests::remove_min_amount_a(), + BalanceForTests::remove_min_amount_b(), + ); +} + #[test] fn test_call_remove_liquidity_chained_call_successful() { let (post_states, chained_calls) = remove_liquidity( @@ -1792,6 +2122,37 @@ fn test_call_remove_liquidity_chained_call_successful() { assert!(chained_call_lp == ChainedCallForTests::cc_remove_pool_lp()); } +#[test] +fn test_call_remove_liquidity_chained_call_with_vault_surplus_successful() { + let (post_states, chained_calls) = remove_liquidity( + AccountWithMetadataForTests::pool_definition_init(), + AccountWithMetadataForTests::vault_a_with_surplus(), + AccountWithMetadataForTests::vault_b_with_surplus(), + AccountWithMetadataForTests::pool_lp_init(), + AccountWithMetadataForTests::user_holding_a(), + AccountWithMetadataForTests::user_holding_b(), + AccountWithMetadataForTests::user_holding_lp_init(), + NonZero::new(BalanceForTests::remove_amount_lp()).unwrap(), + BalanceForTests::remove_min_amount_a(), + BalanceForTests::remove_min_amount_b_surplus(), + ); + + let pool_post = post_states[0].clone(); + + assert!( + AccountWithMetadataForTests::pool_definition_remove_successful().account + == *pool_post.account() + ); + + let chained_call_lp = chained_calls[0].clone(); + let chained_call_b = chained_calls[1].clone(); + let chained_call_a = chained_calls[2].clone(); + + assert!(chained_call_a == ChainedCallForTests::cc_remove_token_a_with_surplus()); + assert!(chained_call_b == ChainedCallForTests::cc_remove_token_b_with_surplus()); + assert!(chained_call_lp == ChainedCallForTests::cc_remove_pool_lp()); +} + #[should_panic(expected = "Balances must be nonzero")] #[test] fn test_call_new_definition_with_zero_balance_1() { @@ -2140,6 +2501,36 @@ fn test_call_swap_below_min_out() { ); } +#[should_panic(expected = "Effective swap amount should be nonzero")] +#[test] +fn test_call_swap_effective_amount_zero() { + let _post_states = swap_exact_input( + AccountWithMetadataForTests::pool_definition_init(), + AccountWithMetadataForTests::vault_a_init(), + AccountWithMetadataForTests::vault_b_init(), + AccountWithMetadataForTests::user_holding_a(), + AccountWithMetadataForTests::user_holding_b(), + 1, + 0, + IdForTests::token_a_definition_id(), + ); +} + +#[should_panic(expected = "Withdraw amount should be nonzero")] +#[test] +fn test_call_swap_output_rounds_to_zero() { + let _post_states = swap_exact_input( + AccountWithMetadataForTests::pool_definition_init_low_balances(), + AccountWithMetadataForTests::vault_a_init_low(), + AccountWithMetadataForTests::vault_b_init_low(), + AccountWithMetadataForTests::user_holding_a(), + AccountWithMetadataForTests::user_holding_b(), + 2, + 0, + IdForTests::token_a_definition_id(), + ); +} + #[test] fn test_call_swap_chained_call_successful_1() { let (post_states, chained_calls) = swap_exact_input( @@ -2755,7 +3146,7 @@ fn test_sync_reserves_panics_when_vault_b_under_collateralized() { } #[test] -fn test_donation_then_add_liquidity_sync_mitigates_mispricing() { +fn test_donation_then_add_liquidity_matches_synced_pricing() { let donation_a = 100u128; let mut donated_vault_a = AccountWithMetadataForTests::vault_a_init(); @@ -2813,7 +3204,7 @@ fn test_donation_then_add_liquidity_sync_mitigates_mispricing() { .unwrap() .liquidity_pool_supply; - assert!(synced_delta_lp < unsynced_delta_lp); + assert_eq!(synced_delta_lp, unsynced_delta_lp); } #[should_panic(expected = "token_a * token_b overflows u128")] @@ -2837,7 +3228,7 @@ fn new_definition_overflow_protection() { ); } -#[should_panic(expected = "reserve_a * max_amount_b overflows u128")] +#[should_panic(expected = "vault_a_balance * max_amount_to_add_token_b overflows u128")] #[test] fn add_liquidity_overflow_protection() { let large_reserve: u128 = u128::MAX / 2 + 1; @@ -2903,7 +3294,7 @@ fn add_liquidity_overflow_protection() { AccountWithMetadataForTests::user_holding_lp_init(), NonZero::new(1).unwrap(), 500, - 2, // max_amount_b=2 → reserve_a * 2 overflows + 2, // max_amount_to_add_token_b=2 → vault_a_balance * 2 overflows ); } @@ -2992,7 +3383,7 @@ fn remove_liquidity_overflow_protection() { ); } -#[should_panic(expected = "reserve * amount_in overflows u128")] +#[should_panic(expected = "reserve_withdraw_vault_amount * effective_amount_in overflows u128")] #[test] fn swap_exact_input_overflow_protection() { let large_reserve: u128 = u128::MAX / 2 + 1; @@ -3048,15 +3439,15 @@ fn swap_exact_input_overflow_protection() { account_id: IdForTests::vault_b_id(), }; - // Swap token_a in: withdraw_amount = reserve_b * swap_amount_in / (reserve_a + swap_amount_in) - // reserve_b is large, so reserve_b * 2 overflows + // With fees applied, swap_amount_in=3 still rounds down to effective_amount_in=2. + // reserve_b is large, so reserve_b * effective_amount_in overflows. let _result = swap_exact_input( pool, vault_a, vault_b, AccountWithMetadataForTests::user_holding_a(), AccountWithMetadataForTests::user_holding_b(), - 2, + 3, 1, IdForTests::token_a_definition_id(), ); diff --git a/integration_tests/tests/amm.rs b/integration_tests/tests/amm.rs index b36da68..6408752 100644 --- a/integration_tests/tests/amm.rs +++ b/integration_tests/tests/amm.rs @@ -144,6 +144,14 @@ impl Balances { 500 } + fn remove_min_a_with_surplus() -> u128 { + 1_100 + } + + fn remove_min_b_with_surplus() -> u128 { + 550 + } + fn add_min_lp() -> u128 { 1_000 } @@ -164,8 +172,16 @@ impl Balances { 200 } + fn reserve_a_swap_1() -> u128 { + 3_575 + } + + fn reserve_b_swap_1() -> u128 { + 3_497 + } + fn vault_a_swap_1() -> u128 { - 3_572 + 3_575 } fn vault_b_swap_1() -> u128 { @@ -173,19 +189,27 @@ impl Balances { } fn user_a_swap_1() -> u128 { - 11_428 + 11_425 } fn user_b_swap_1() -> u128 { 9_000 } + fn reserve_a_swap_2() -> u128 { + 5_997 + } + + fn reserve_b_swap_2() -> u128 { + 2_085 + } + fn vault_a_swap_2() -> u128 { 6_000 } fn vault_b_swap_2() -> u128 { - 2_084 + 2_085 } fn user_a_swap_2() -> u128 { @@ -193,7 +217,7 @@ impl Balances { } fn user_b_swap_2() -> u128 { - 10_416 + 10_415 } fn vault_a_add() -> u128 { @@ -228,6 +252,22 @@ impl Balances { 2_000 } + fn vault_a_with_surplus() -> u128 { + 5_500 + } + + fn vault_b_with_surplus() -> u128 { + 2_750 + } + + fn vault_a_remove_with_surplus() -> u128 { + 4_400 + } + + fn vault_b_remove_with_surplus() -> u128 { + 2_200 + } + fn user_a_remove() -> u128 { 11_000 } @@ -236,6 +276,14 @@ impl Balances { 10_500 } + fn user_a_remove_with_surplus() -> u128 { + 11_100 + } + + fn user_b_remove_with_surplus() -> u128 { + 10_550 + } + fn user_lp_remove() -> u128 { 1_000 } @@ -369,6 +417,30 @@ impl Accounts { } } + fn vault_a_with_surplus() -> Account { + Account { + program_owner: Ids::token_program(), + balance: 0_u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: Ids::token_a_definition(), + balance: Balances::vault_a_with_surplus(), + }), + nonce: Nonce(0), + } + } + + fn vault_b_with_surplus() -> Account { + Account { + program_owner: Ids::token_program(), + balance: 0_u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: Ids::token_b_definition(), + balance: Balances::vault_b_with_surplus(), + }), + nonce: Nonce(0), + } + } + fn user_lp_holding() -> Account { Account { program_owner: Ids::token_program(), @@ -406,8 +478,8 @@ impl Accounts { vault_b_id: Ids::vault_b(), liquidity_pool_id: Ids::token_lp_definition(), liquidity_pool_supply: Balances::pool_lp_supply_init(), - reserve_a: Balances::vault_a_swap_1(), - reserve_b: Balances::vault_b_swap_1(), + reserve_a: Balances::reserve_a_swap_1(), + reserve_b: Balances::reserve_b_swap_1(), fees: Balances::fee_tier(), active: true, }), @@ -474,8 +546,8 @@ impl Accounts { vault_b_id: Ids::vault_b(), liquidity_pool_id: Ids::token_lp_definition(), liquidity_pool_supply: Balances::pool_lp_supply_init(), - reserve_a: Balances::vault_a_swap_2(), - reserve_b: Balances::vault_b_swap_2(), + reserve_a: Balances::reserve_a_swap_2(), + reserve_b: Balances::reserve_b_swap_2(), fees: Balances::fee_tier(), active: true, }), @@ -668,6 +740,30 @@ impl Accounts { } } + fn vault_a_remove_with_surplus() -> Account { + Account { + program_owner: Ids::token_program(), + balance: 0_u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: Ids::token_a_definition(), + balance: Balances::vault_a_remove_with_surplus(), + }), + nonce: Nonce(0), + } + } + + fn vault_b_remove_with_surplus() -> Account { + Account { + program_owner: Ids::token_program(), + balance: 0_u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: Ids::token_b_definition(), + balance: Balances::vault_b_remove_with_surplus(), + }), + nonce: Nonce(0), + } + } + fn user_a_holding_remove() -> Account { Account { program_owner: Ids::token_program(), @@ -692,6 +788,30 @@ impl Accounts { } } + fn user_a_holding_remove_with_surplus() -> Account { + Account { + program_owner: Ids::token_program(), + balance: 0_u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: Ids::token_a_definition(), + balance: Balances::user_a_remove_with_surplus(), + }), + nonce: Nonce(0), + } + } + + fn user_b_holding_remove_with_surplus() -> Account { + Account { + program_owner: Ids::token_program(), + balance: 0_u128, + data: Data::from(&TokenHolding::Fungible { + definition_id: Ids::token_b_definition(), + balance: Balances::user_b_remove_with_surplus(), + }), + nonce: Nonce(0), + } + } + fn user_lp_holding_remove() -> Account { Account { program_owner: Ids::token_program(), @@ -925,6 +1045,10 @@ fn state_for_amm_tests_with_new_def() -> V03State { state } +fn current_nonce(state: &V03State, account_id: AccountId) -> Nonce { + state.get_account_by_id(account_id).nonce +} + fn try_execute_new_definition(state: &mut V03State, fees: u128) -> Result<(), NssaError> { let instruction = amm_core::Instruction::NewDefinition { token_a_amount: Balances::vault_a_init(), @@ -945,7 +1069,10 @@ fn try_execute_new_definition(state: &mut V03State, fees: u128) -> Result<(), Ns Ids::user_b(), Ids::user_lp(), ], - vec![Nonce(0), Nonce(0)], + vec![ + current_nonce(state, Ids::user_a()), + current_nonce(state, Ids::user_b()), + ], instruction, ) .unwrap(); @@ -961,6 +1088,163 @@ fn execute_new_definition(state: &mut V03State, fees: u128) { try_execute_new_definition(state, fees).unwrap(); } +fn execute_swap_a_to_b(state: &mut V03State, swap_amount_in: u128, min_amount_out: u128) { + let instruction = amm_core::Instruction::SwapExactInput { + swap_amount_in, + min_amount_out, + token_definition_id_in: Ids::token_a_definition(), + }; + + let message = public_transaction::Message::try_new( + Ids::amm_program(), + vec![ + Ids::pool_definition(), + Ids::vault_a(), + Ids::vault_b(), + Ids::user_a(), + Ids::user_b(), + ], + vec![current_nonce(state, Ids::user_a())], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::user_a()]); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx, 0).unwrap(); +} + +fn execute_swap_b_to_a(state: &mut V03State, swap_amount_in: u128, min_amount_out: u128) { + let instruction = amm_core::Instruction::SwapExactInput { + swap_amount_in, + min_amount_out, + token_definition_id_in: Ids::token_b_definition(), + }; + + let message = public_transaction::Message::try_new( + Ids::amm_program(), + vec![ + Ids::pool_definition(), + Ids::vault_a(), + Ids::vault_b(), + Ids::user_a(), + Ids::user_b(), + ], + vec![current_nonce(state, Ids::user_b())], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::user_b()]); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx, 0).unwrap(); +} + +fn execute_add_liquidity( + state: &mut V03State, + min_amount_liquidity: u128, + max_amount_to_add_token_a: u128, + max_amount_to_add_token_b: u128, +) { + let instruction = amm_core::Instruction::AddLiquidity { + min_amount_liquidity, + max_amount_to_add_token_a, + max_amount_to_add_token_b, + }; + + let message = public_transaction::Message::try_new( + Ids::amm_program(), + vec![ + Ids::pool_definition(), + Ids::vault_a(), + Ids::vault_b(), + Ids::token_lp_definition(), + Ids::user_a(), + Ids::user_b(), + Ids::user_lp(), + ], + vec![ + current_nonce(state, Ids::user_a()), + current_nonce(state, Ids::user_b()), + ], + instruction, + ) + .unwrap(); + + let witness_set = + public_transaction::WitnessSet::for_message(&message, &[&Keys::user_a(), &Keys::user_b()]); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx, 0).unwrap(); +} + +fn execute_remove_liquidity( + state: &mut V03State, + remove_liquidity_amount: u128, + min_amount_to_remove_token_a: u128, + min_amount_to_remove_token_b: u128, +) { + let instruction = amm_core::Instruction::RemoveLiquidity { + remove_liquidity_amount, + min_amount_to_remove_token_a, + min_amount_to_remove_token_b, + }; + + let message = public_transaction::Message::try_new( + Ids::amm_program(), + vec![ + Ids::pool_definition(), + Ids::vault_a(), + Ids::vault_b(), + Ids::token_lp_definition(), + Ids::user_a(), + Ids::user_b(), + Ids::user_lp(), + ], + vec![current_nonce(state, Ids::user_lp())], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::user_lp()]); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx, 0).unwrap(); +} + +fn fungible_balance(account: &Account) -> u128 { + let holding = TokenHolding::try_from(&account.data).expect("expected token holding"); + let TokenHolding::Fungible { + definition_id: _, + balance, + } = holding + else { + panic!("expected fungible token holding") + }; + + balance +} + +fn pool_definition(account: &Account) -> PoolDefinition { + PoolDefinition::try_from(&account.data).expect("expected pool definition") +} + +fn fungible_total_supply(account: &Account) -> u128 { + let definition = TokenDefinition::try_from(&account.data).expect("expected token definition"); + let TokenDefinition::Fungible { + name: _, + total_supply, + metadata_id: _, + } = definition + else { + panic!("expected fungible token definition") + }; + + total_supply +} + #[test] fn amm_remove_liquidity() { let mut state = state_for_amm_tests(); @@ -1055,6 +1339,69 @@ fn amm_remove_liquidity_insufficient_user_lp_fails() { assert!(state.transition_from_public_transaction(&tx, 0).is_err()); } +#[test] +fn amm_remove_liquidity_with_surplus() { + let mut state = state_for_amm_tests(); + state.force_insert_account(Ids::vault_a(), Accounts::vault_a_with_surplus()); + state.force_insert_account(Ids::vault_b(), Accounts::vault_b_with_surplus()); + + let instruction = amm_core::Instruction::RemoveLiquidity { + remove_liquidity_amount: Balances::remove_lp(), + min_amount_to_remove_token_a: Balances::remove_min_a_with_surplus(), + min_amount_to_remove_token_b: Balances::remove_min_b_with_surplus(), + }; + + let message = public_transaction::Message::try_new( + Ids::amm_program(), + vec![ + Ids::pool_definition(), + Ids::vault_a(), + Ids::vault_b(), + Ids::token_lp_definition(), + Ids::user_a(), + Ids::user_b(), + Ids::user_lp(), + ], + vec![Nonce(0)], + instruction, + ) + .unwrap(); + + let witness_set = public_transaction::WitnessSet::for_message(&message, &[&Keys::user_lp()]); + + let tx = PublicTransaction::new(message, witness_set); + state.transition_from_public_transaction(&tx, 0).unwrap(); + + assert_eq!( + state.get_account_by_id(Ids::pool_definition()), + Accounts::pool_definition_remove() + ); + assert_eq!( + state.get_account_by_id(Ids::vault_a()), + Accounts::vault_a_remove_with_surplus() + ); + assert_eq!( + state.get_account_by_id(Ids::vault_b()), + Accounts::vault_b_remove_with_surplus() + ); + assert_eq!( + state.get_account_by_id(Ids::token_lp_definition()), + Accounts::token_lp_definition_remove() + ); + assert_eq!( + state.get_account_by_id(Ids::user_a()), + Accounts::user_a_holding_remove_with_surplus() + ); + assert_eq!( + state.get_account_by_id(Ids::user_b()), + Accounts::user_b_holding_remove_with_surplus() + ); + assert_eq!( + state.get_account_by_id(Ids::user_lp()), + Accounts::user_lp_holding_remove() + ); +} + #[test] fn amm_new_definition_inactive_initialized_pool_and_uninit_user_lp() { let mut state = state_for_amm_tests_with_new_def(); @@ -1421,3 +1768,106 @@ fn amm_swap_a_to_b() { Accounts::user_b_holding_swap_2() ); } + +#[test] +fn amm_fee_accumulates_across_multiple_swaps_and_pays_out_on_remove() { + let mut state = state_for_amm_tests(); + + execute_swap_a_to_b(&mut state, 1_000, 200); + execute_swap_b_to_a(&mut state, 1_000, 200); + + let pool_before_remove = pool_definition(&state.get_account_by_id(Ids::pool_definition())); + assert_eq!(pool_before_remove.reserve_a, 4_058); + assert_eq!(pool_before_remove.reserve_b, 3_082); + assert_eq!(pool_before_remove.fees, Balances::fee_tier()); + + let vault_a_before_remove = fungible_balance(&state.get_account_by_id(Ids::vault_a())); + let vault_b_before_remove = fungible_balance(&state.get_account_by_id(Ids::vault_b())); + assert_eq!(vault_a_before_remove, 4_061); + assert_eq!(vault_b_before_remove, 3_085); + assert_eq!(vault_a_before_remove - pool_before_remove.reserve_a, 3); + assert_eq!(vault_b_before_remove - pool_before_remove.reserve_b, 3); + + execute_remove_liquidity(&mut state, 1_000, 812, 617); + + let pool_after_remove = pool_definition(&state.get_account_by_id(Ids::pool_definition())); + assert_eq!(pool_after_remove.reserve_a, 3_247); + assert_eq!(pool_after_remove.reserve_b, 2_466); + assert_eq!(pool_after_remove.liquidity_pool_supply, 4_000); + + let vault_a_after_remove = fungible_balance(&state.get_account_by_id(Ids::vault_a())); + let vault_b_after_remove = fungible_balance(&state.get_account_by_id(Ids::vault_b())); + assert_eq!(vault_a_after_remove, 3_249); + assert_eq!(vault_b_after_remove, 2_468); + assert_eq!(vault_a_after_remove - pool_after_remove.reserve_a, 2); + assert_eq!(vault_b_after_remove - pool_after_remove.reserve_b, 2); + + assert_eq!( + fungible_balance(&state.get_account_by_id(Ids::user_a())), + 11_751 + ); + assert_eq!( + fungible_balance(&state.get_account_by_id(Ids::user_b())), + 10_032 + ); + assert_eq!( + fungible_balance(&state.get_account_by_id(Ids::user_lp())), + 1_000 + ); + assert_eq!( + fungible_total_supply(&state.get_account_by_id(Ids::token_lp_definition())), + 4_000 + ); +} + +#[test] +fn amm_add_liquidity_after_fee_accrual_preserves_surplus() { + let mut state = state_for_amm_tests(); + + execute_swap_a_to_b(&mut state, 1_000, 200); + execute_swap_b_to_a(&mut state, 1_000, 200); + execute_swap_a_to_b(&mut state, 1_000, 200); + execute_swap_b_to_a(&mut state, 1_000, 200); + + let pool_before_add = pool_definition(&state.get_account_by_id(Ids::pool_definition())); + let vault_a_before_add = fungible_balance(&state.get_account_by_id(Ids::vault_a())); + let vault_b_before_add = fungible_balance(&state.get_account_by_id(Ids::vault_b())); + + assert_eq!(pool_before_add.reserve_a, 3_604); + assert_eq!(pool_before_add.reserve_b, 3_472); + assert_eq!(vault_a_before_add, 3_610); + assert_eq!(vault_b_before_add, 3_478); + assert_eq!(vault_a_before_add - pool_before_add.reserve_a, 6); + assert_eq!(vault_b_before_add - pool_before_add.reserve_b, 6); + + execute_add_liquidity(&mut state, 1_436, 2_000, 1_000); + + let pool_after_add = pool_definition(&state.get_account_by_id(Ids::pool_definition())); + let vault_a_after_add = fungible_balance(&state.get_account_by_id(Ids::vault_a())); + let vault_b_after_add = fungible_balance(&state.get_account_by_id(Ids::vault_b())); + + assert_eq!(pool_after_add.reserve_a, 4_641); + assert_eq!(pool_after_add.reserve_b, 4_472); + assert_eq!(pool_after_add.liquidity_pool_supply, 6_436); + assert_eq!(vault_a_after_add, 4_647); + assert_eq!(vault_b_after_add, 4_478); + assert_eq!(vault_a_after_add - pool_after_add.reserve_a, 6); + assert_eq!(vault_b_after_add - pool_after_add.reserve_b, 6); + + assert_eq!( + fungible_balance(&state.get_account_by_id(Ids::user_a())), + 10_353 + ); + assert_eq!( + fungible_balance(&state.get_account_by_id(Ids::user_b())), + 8_022 + ); + assert_eq!( + fungible_balance(&state.get_account_by_id(Ids::user_lp())), + 3_436 + ); + assert_eq!( + fungible_total_supply(&state.get_account_by_id(Ids::token_lp_definition())), + 6_436 + ); +}