diff --git a/programs/omnipair/src/instructions/lending/flashloan.rs b/programs/omnipair/src/instructions/lending/flashloan.rs index e91d0d4..6d5ef15 100644 --- a/programs/omnipair/src/instructions/lending/flashloan.rs +++ b/programs/omnipair/src/instructions/lending/flashloan.rs @@ -281,7 +281,7 @@ impl<'info> Flashloan<'info> { let required_balance0 = balance0_before.checked_add(fee0).unwrap(); let required_balance1 = balance1_before.checked_add(fee1).unwrap(); - + require!( token0_vault.amount >= required_balance0, ErrorCode::InsufficientAmount0 @@ -291,6 +291,14 @@ impl<'info> Flashloan<'info> { ErrorCode::InsufficientAmount1 ); + // Calculate excess repaid tokens (if any) and accumulate them + let excess0 = token0_vault.amount.checked_sub(required_balance0).unwrap_or(0); + let excess1 = token1_vault.amount.checked_sub(required_balance1).unwrap_or(0); + + // update the reserves: fee + excess repaid + pair.reserve0 = pair.reserve0.checked_add(fee0).unwrap().checked_add(excess0).unwrap(); + pair.reserve1 = pair.reserve1.checked_add(fee1).unwrap().checked_add(excess1).unwrap(); + // Emit event emit_cpi!(FlashloanEvent { amount0, diff --git a/programs/omnipair/src/instructions/spot/swap.rs b/programs/omnipair/src/instructions/spot/swap.rs index 8bf1d37..7d23eb8 100644 --- a/programs/omnipair/src/instructions/spot/swap.rs +++ b/programs/omnipair/src/instructions/spot/swap.rs @@ -159,15 +159,15 @@ impl<'info> Swap<'info> { let last_k = (pair.reserve0 as u128).checked_mul(pair.reserve1 as u128).ok_or(ErrorCode::InvariantOverflow)?; let is_token0_in = user_token_in_account.mint == pair.token0; - // Calculate total fee amount - let total_fee = (amount_in as u128) + // Swap fee = LP fee + Futarchy fee + let swap_fee = (amount_in as u128) .checked_mul(pair.swap_fee_bps as u128) .ok_or(ErrorCode::FeeMathOverflow)? .checked_div(BPS_DENOMINATOR as u128) .ok_or(ErrorCode::FeeMathOverflow)? as u64; - // Calculate futarchy fee portion of the total fee - let futarchy_fee = (total_fee as u128) + // Calculate futarchy fee portion of the swap fee + let futarchy_fee = (swap_fee as u128) .checked_mul(futarchy_authority.revenue_share.swap_bps as u128) .ok_or(ErrorCode::FeeMathOverflow)? .checked_div(BPS_DENOMINATOR as u128) @@ -190,23 +190,17 @@ impl<'info> Swap<'info> { )?; } - // amount_in_after_fee = amount_in * (10000 - swap_fee_bps) / 10000 - let amount_in_after_fee = (amount_in as u128) - .checked_mul((BPS_DENOMINATOR as u128).checked_sub(pair.swap_fee_bps as u128).ok_or(ErrorCode::FeeMathOverflow)?) - .ok_or(ErrorCode::FeeMathOverflow)? - .checked_div(BPS_DENOMINATOR as u128) - .ok_or(ErrorCode::FeeMathOverflow)? - .try_into() - .map_err(|_| ErrorCode::FeeMathOverflow)?; + // amount_in_after_swap_fee = amount_in - swap_fee + let amount_in_after_swap_fee = amount_in.checked_sub(swap_fee).ok_or(ErrorCode::FeeMathOverflow)?; let reserve_in = if is_token0_in { pair.reserve0 } else { pair.reserve1 }; let reserve_out = if is_token0_in { pair.reserve1 } else { pair.reserve0 }; // Δy = (Δx * y) / (x + Δx) let denominator = (reserve_in as u128) - .checked_add(amount_in_after_fee as u128) + .checked_add(amount_in_after_swap_fee as u128) .ok_or(ErrorCode::DenominatorOverflow)?; - let amount_out = (amount_in_after_fee as u128) + let amount_out = (amount_in_after_swap_fee as u128) .checked_mul(reserve_out as u128) .ok_or(ErrorCode::OutputAmountOverflow)? .checked_div(denominator) @@ -214,7 +208,10 @@ impl<'info> Swap<'info> { .try_into() .map_err(|_| ErrorCode::OutputAmountOverflow)?; - let new_reserve_in = reserve_in.checked_add(amount_in_after_fee).ok_or(ErrorCode::Overflow)?; + // Calculate the amount in with the LP portion of the fee: + // amount_in_with_lp_fee = amount_in - swap_fee + lp_fee = amount_in - futarchy_fee + let amount_in_with_lp_fee = amount_in.checked_sub(futarchy_fee).ok_or(ErrorCode::Overflow)?; + let new_reserve_in = reserve_in.checked_add(amount_in_with_lp_fee).ok_or(ErrorCode::Overflow)?; let new_reserve_out = reserve_out.checked_sub(amount_out).ok_or(ErrorCode::Overflow)?; require_gte!(amount_out, min_amount_out, ErrorCode::InsufficientOutputAmount); @@ -268,7 +265,7 @@ impl<'info> Swap<'info> { is_token0_in, amount_in: amount_in, amount_out: amount_out, - amount_in_after_fee: amount_in_after_fee, + amount_in_after_fee: amount_in_after_swap_fee as u64, }); Ok(())