refactor: improve AMM error handling and mathematical calculations

franRappazzini · franRappazzini · commit ea515ec03ef9 · 2026-01-21T14:04:31.000-03:00
- Enhanced error handling in swap and deposit liquidity instructions
- Improved mathematical utilities for precision and safety
- Updated test suites to cover edge cases with decimals and slippage
- Refined type definitions in fuzz tests
diff --git a/programs/amm/src/errors.rs b/programs/amm/src/errors.rs
@@ -26,6 +26,4 @@ pub enum AmmError {
     DecimalsTooLarge,
     #[msg("Fee exceeds maximum allowed")]
     FeeExceedsMaximum,
-    #[msg("Slippage tolerance exceeded")]
-    SlippageExceeded,
 }
diff --git a/programs/amm/src/instructions/deposit_liquidity.rs b/programs/amm/src/instructions/deposit_liquidity.rs
@@ -137,7 +137,10 @@ impl<'info> DepositLiquidity<'info> {
             self.liquidity_pool.amount_mint_b,
         )?;
 
-        require!(liquidity >= min_lp_out, AmmError::SlippageExceeded);
+        require!(
+            liquidity >= min_lp_out,
+            AmmError::InsufficientLiquidityMinted
+        );
 
         let signer_seeds: &[&[&[u8]]] = &[&[
             LIQUIDITY_POOL_SEED,
diff --git a/programs/amm/src/instructions/swap.rs b/programs/amm/src/instructions/swap.rs
@@ -65,19 +65,13 @@ pub struct Swap<'info> {
 }
 
 impl<'info> Swap<'info> {
-    pub fn swap(&mut self, pool_id: u64, params: SwapParams, slippage_limit: u64) -> Result<()> {
+    pub fn swap(&mut self, pool_id: u64, params: SwapParams) -> Result<()> {
         // calculate and transfer input mint from signer to vault
         let input_is_mint_a = self.liquidity_pool.mint_a == self.input_mint.key();
 
         let (input_amount, output_amount, protocol_fee_amount) =
             self.calculate_amounts(&params, input_is_mint_a)?;
 
-        if params.is_exact_in() {
-            require!(output_amount >= slippage_limit, AmmError::SlippageExceeded);
-        } else if slippage_limit > 0 {
-            require!(input_amount <= slippage_limit, AmmError::SlippageExceeded);
-        }
-
         msg!(
             "Swapping {} of input mint for {} of output mint",
             input_amount,
@@ -175,8 +169,14 @@ impl<'info> Swap<'info> {
 
 #[derive(AnchorSerialize, AnchorDeserialize)]
 pub enum SwapParams {
-    ExactIn { input_amount: u64 },
-    ExactOut { output_amount: u64 },
+    ExactIn {
+        input_amount: u64,
+        min_output_amount: u64,
+    },
+    ExactOut {
+        output_amount: u64,
+        max_input_amount: u64,
+    },
 }
 
 impl SwapParams {
diff --git a/programs/amm/src/lib.rs b/programs/amm/src/lib.rs
@@ -34,15 +34,16 @@ pub mod amm {
         amount_b: u64,
         min_lp_out: u64,
     ) -> Result<()> {
-        ctx.accounts.deposit_liquidity(pool_id, amount_a, amount_b, min_lp_out)
+        ctx.accounts
+            .deposit_liquidity(pool_id, amount_a, amount_b, min_lp_out)
     }
 
     pub fn redeem_lp(ctx: Context<RedeemLp>, pool_id: u64, lp_amount: u64) -> Result<()> {
         ctx.accounts.redeem_lp(pool_id, lp_amount)
     }
 
-    pub fn swap(ctx: Context<Swap>, pool_id: u64, params: SwapParams, slippage_limit: u64) -> Result<()> {
-        ctx.accounts.swap(pool_id, params, slippage_limit)
+    pub fn swap(ctx: Context<Swap>, pool_id: u64, params: SwapParams) -> Result<()> {
+        ctx.accounts.swap(pool_id, params)
     }
 
     pub fn withdraw_protocol_fees(ctx: Context<WithdrawProtocolFees>, pool_id: u64) -> Result<()> {
diff --git a/programs/amm/src/utils/math.rs b/programs/amm/src/utils/math.rs
@@ -314,7 +314,10 @@ pub fn calculate_swap_amounts(
         .ok_or(AmmError::MathUnderflow)?;
 
     match swap_params {
-        SwapParams::ExactIn { input_amount } => {
+        SwapParams::ExactIn {
+            input_amount,
+            min_output_amount,
+        } => {
             require!(*input_amount > 0, AmmError::InsufficientInputAmount);
             require!(
                 *input_amount < reserve_input,
@@ -345,6 +348,12 @@ pub fn calculate_swap_amounts(
                 .try_into()
                 .map_err(|_| AmmError::MathOverflow)?;
 
+            // slippage check
+            require!(
+                output_amount >= *min_output_amount,
+                AmmError::InsufficientOutputAmount
+            );
+
             // Calcular el fee del protocolo del input_amount
             let protocol_fee_amount = (*input_amount as u128)
                 .checked_mul(protocol_fee_bps as u128)
@@ -356,7 +365,10 @@ pub fn calculate_swap_amounts(
 
             Ok((*input_amount, output_amount, protocol_fee_amount))
         }
-        SwapParams::ExactOut { output_amount } => {
+        SwapParams::ExactOut {
+            output_amount,
+            max_input_amount,
+        } => {
             require!(*output_amount > 0, AmmError::InsufficientOutputAmount);
             require!(
                 *output_amount < reserve_output,
@@ -383,6 +395,14 @@ pub fn calculate_swap_amounts(
                 .try_into()
                 .map_err(|_| AmmError::MathOverflow)?;
 
+            // slippage check
+            if *max_input_amount > 0 {
+                require!(
+                    input_amount <= *max_input_amount,
+                    AmmError::ExcessiveInputAmount
+                );
+            }
+
             // Calcular el fee del protocolo del input_amount
             let protocol_fee_amount = (input_amount as u128)
                 .checked_mul(protocol_fee_bps as u128)
diff --git a/tests/amm-decimals-local.test.ts b/tests/amm-decimals-local.test.ts
@@ -445,9 +445,14 @@ describe("amm - decimals tests", () => {
       const [poolBefore] = await getLiquidityPoolAccount(poolId);
 
       const inputAmount = 10_000_000; // 10 USDC
+      const minOutputAmount = 0;
+
+      const param: anchor.IdlTypes<Amm>["swapParams"] = {
+        exactIn: { inputAmount: bn(inputAmount), minOutputAmount: bn(minOutputAmount) },
+      };
 
       const tx = await program.methods
-        .swap(bn(poolId), { exactIn: { inputAmount: bn(inputAmount) } }, bn(0))
+        .swap(bn(poolId), param)
         .accounts({
           tokenProgram: TOKEN_PROGRAM_ID,
           signer: liquidityProvider.publicKey,
@@ -469,9 +474,14 @@ describe("amm - decimals tests", () => {
       const [poolBefore] = await getLiquidityPoolAccount(poolId);
 
       const inputAmount = 100_000_000; // 0.1 SOL
+      const minOutputAmount = 0;
+
+      const param: anchor.IdlTypes<Amm>["swapParams"] = {
+        exactIn: { inputAmount: bn(inputAmount), minOutputAmount: bn(minOutputAmount) },
+      };
 
       const tx = await program.methods
-        .swap(bn(poolId), { exactIn: { inputAmount: bn(inputAmount) } }, bn(0))
+        .swap(bn(poolId), param)
         .accounts({
           tokenProgram: TOKEN_PROGRAM_ID,
           signer: liquidityProvider.publicKey,
diff --git a/tests/amm-local.test.ts b/tests/amm-local.test.ts
@@ -349,18 +349,18 @@ describe("amm", () => {
 
   it("`swap` method using 'exact in' param with mint A!", async () => {
     const INPUT_AMOUNT = 10_000_000; // 10 token
-    const SLIPPAGE_LIMIT = 0; // expressed in amount, no bps
+    const MIN_OUTPUT_AMOUNT = 0;
 
     const param: anchor.IdlTypes<Amm>["swapParams"] = {
-      exactIn: { inputAmount: bn(INPUT_AMOUNT) },
+      exactIn: { inputAmount: bn(INPUT_AMOUNT), minOutputAmount: bn(MIN_OUTPUT_AMOUNT) },
     };
 
     const [prevLiquidityPool] = await getLiquidityPoolAccount(poolId);
     console.log("A amount in vault:", prevLiquidityPool.amountMintA);
     console.log("B amount in vault:", prevLiquidityPool.amountMintB);
 
     const tx = await program.methods
-      .swap(bn(poolId), param, bn(SLIPPAGE_LIMIT))
+      .swap(bn(poolId), param)
       .accounts({
         signer: user.publicKey,
         inputMint: mintA,
@@ -397,18 +397,18 @@ describe("amm", () => {
 
   it("`swap` method using 'exact out' param with mint A!", async () => {
     const OUTPUT_AMOUNT = 20_000_000; // 2 token
-    const SLIPPAGE_LIMIT = 0; // expressed in amount, no bps
+    const MAX_INPUT_AMOUNT = 0;
 
     const param: anchor.IdlTypes<Amm>["swapParams"] = {
-      exactOut: { outputAmount: bn(OUTPUT_AMOUNT) },
+      exactOut: { outputAmount: bn(OUTPUT_AMOUNT), maxInputAmount: bn(MAX_INPUT_AMOUNT) },
     };
 
     const [prevLiquidityPool] = await getLiquidityPoolAccount(poolId);
     console.log("A amount in vault:", prevLiquidityPool.amountMintA);
     console.log("B amount in vault:", prevLiquidityPool.amountMintB);
 
     const tx = await program.methods
-      .swap(bn(poolId), param, bn(SLIPPAGE_LIMIT))
+      .swap(bn(poolId), param)
       .accounts({
         signer: user.publicKey,
         inputMint: mintA,
@@ -445,18 +445,18 @@ describe("amm", () => {
 
   it("`swap` method using 'exact in' param with mint B!", async () => {
     const INPUT_AMOUNT = 10_000_000; // 10 token
-    const SLIPPAGE_LIMIT = 0; // expressed in amount, no bps
+    const MIN_OUTPUT_AMOUNT = 0;
 
     const param: anchor.IdlTypes<Amm>["swapParams"] = {
-      exactIn: { inputAmount: bn(INPUT_AMOUNT) },
+      exactIn: { inputAmount: bn(INPUT_AMOUNT), minOutputAmount: bn(MIN_OUTPUT_AMOUNT) },
     };
 
     const [prevLiquidityPool] = await getLiquidityPoolAccount(poolId);
     console.log("A amount in vault:", prevLiquidityPool.amountMintA);
     console.log("B amount in vault:", prevLiquidityPool.amountMintB);
 
     const tx = await program.methods
-      .swap(bn(poolId), param, bn(SLIPPAGE_LIMIT))
+      .swap(bn(poolId), param)
       .accounts({
         signer: user.publicKey,
         inputMint: mintB,
@@ -493,18 +493,18 @@ describe("amm", () => {
 
   it("`swap` method using 'exact out' param with mint B!", async () => {
     const OUTPUT_AMOUNT = 2_000_000; // 2 token
-    const SLIPPAGE_LIMIT = 0; // expressed in amount, no bps
+    const MAX_INPUT_AMOUNT = 0;
 
     const param: anchor.IdlTypes<Amm>["swapParams"] = {
-      exactOut: { outputAmount: bn(OUTPUT_AMOUNT) },
+      exactOut: { outputAmount: bn(OUTPUT_AMOUNT), maxInputAmount: bn(MAX_INPUT_AMOUNT) },
     };
 
     const [prevLiquidityPool] = await getLiquidityPoolAccount(poolId);
     console.log("A amount in vault:", prevLiquidityPool.amountMintA);
     console.log("B amount in vault:", prevLiquidityPool.amountMintB);
 
     const tx = await program.methods
-      .swap(bn(poolId), param, bn(SLIPPAGE_LIMIT))
+      .swap(bn(poolId), param)
       .accounts({
         signer: user.publicKey,
         inputMint: mintB,
diff --git a/tests/amm-slippage-local.test.ts b/tests/amm-slippage-local.test.ts
diff --git a/trident-tests/fuzz_0/types.rs b/trident-tests/fuzz_0/types.rs

Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,4 @@ pub enum AmmError {`
`26`	`26`	`DecimalsTooLarge,`
`27`	`27`	`#[msg("Fee exceeds maximum allowed")]`
`28`	`28`	`FeeExceedsMaximum,`
`29`		`- #[msg("Slippage tolerance exceeded")]`
`30`		`- SlippageExceeded,`
`31`	`29`	`}`
Original file line number	Diff line number	Diff line change
`@@ -34,15 +34,16 @@ pub mod amm {`
`34`	`34`	`amount_b: u64,`
`35`	`35`	`min_lp_out: u64,`
`36`	`36`	`) -> Result<()> {`
`37`		`- ctx.accounts.deposit_liquidity(pool_id, amount_a, amount_b, min_lp_out)`
	`37`	`+ ctx.accounts`
	`38`	`+ .deposit_liquidity(pool_id, amount_a, amount_b, min_lp_out)`
`38`	`39`	`}`
`39`	`40`
`40`	`41`	`pub fn redeem_lp(ctx: Context<RedeemLp>, pool_id: u64, lp_amount: u64) -> Result<()> {`
`41`	`42`	`ctx.accounts.redeem_lp(pool_id, lp_amount)`
`42`	`43`	`}`
`43`	`44`
`44`		`- pub fn swap(ctx: Context<Swap>, pool_id: u64, params: SwapParams, slippage_limit: u64) -> Result<()> {`
`45`		`- ctx.accounts.swap(pool_id, params, slippage_limit)`
	`45`	`+ pub fn swap(ctx: Context<Swap>, pool_id: u64, params: SwapParams) -> Result<()> {`
	`46`	`+ ctx.accounts.swap(pool_id, params)`
`46`	`47`	`}`
`47`	`48`
`48`	`49`	`pub fn withdraw_protocol_fees(ctx: Context<WithdrawProtocolFees>, pool_id: u64) -> Result<()> {`