fix deposit ratio bug (#505)

Author: a-hariti

tokens/token-swap/anchor/programs/token-swap/src/instructions/deposit_liquidity.rs

@@ -37,25 +37,30 @@ pub fn deposit_liquidity(
         // Add as is if there is no liquidity
         (amount_a, amount_b)
     } else {
-        let ratio = I64F64::from_num(pool_a.amount)
-            .checked_mul(I64F64::from_num(pool_b.amount))
+        // u128 is enough precision here
+        let amount_a_u128 = amount_a as u128;
+        let amount_b_u128 = amount_b as u128;
+        let pool_a_u128 = pool_a.amount as u128;
+        let pool_b_u128 = pool_b.amount as u128;
+
+        // Calculate the amount of B required if we deposit all of A provided
+        let amount_b_required = amount_a_u128
+            .checked_mul(pool_b_u128)
+            .unwrap()
+            .checked_div(pool_a_u128)
             .unwrap();
-        if pool_a.amount > pool_b.amount {
-            (
-                I64F64::from_num(amount_b)
-                    .checked_mul(ratio)
-                    .unwrap()
-                    .to_num::<u64>(),
-                amount_b,
-            )
+
+        if amount_b_required <= amount_b_u128 {
+            // We have enough B to match the A provided
+            (amount_a, amount_b_required as u64)
         } else {
-            (
-                amount_a,
-                I64F64::from_num(amount_a)
-                    .checked_div(ratio)
-                    .unwrap()
-                    .to_num::<u64>(),
-            )
+            // We don't have enough B, so we must limit by B and calculate A required
+            let amount_a_required = amount_b_u128
+                .checked_mul(pool_a_u128)
+                .unwrap()
+                .checked_div(pool_b_u128)
+                .unwrap();
+            (amount_a_required as u64, amount_b)
         }
     };
 

tokens/token-swap/anchor/tests/deposit-liquidity.ts

@@ -71,4 +71,114 @@ describe("Deposit liquidity", () => {
     const depositTokenAccountB = await connection.getTokenAccountBalance(values.holderAccountB);
     expect(depositTokenAccountB.value.amount).to.equal(values.defaultSupply.sub(values.depositAmountA).toString());
   });
-});
+
+  it('Deposit with existing liquidity (same ratio)', async () => {
+    // 1. Initial Deposit
+    await program.methods
+      .depositLiquidity(values.depositAmountA, values.depositAmountA)
+      .accounts({
+        pool: values.poolKey,
+        poolAuthority: values.poolAuthority,
+        depositor: values.admin.publicKey,
+        mintLiquidity: values.mintLiquidity,
+        mintA: values.mintAKeypair.publicKey,
+        mintB: values.mintBKeypair.publicKey,
+        poolAccountA: values.poolAccountA,
+        poolAccountB: values.poolAccountB,
+        depositorAccountLiquidity: values.liquidityAccount,
+        depositorAccountA: values.holderAccountA,
+        depositorAccountB: values.holderAccountB,
+      })
+      .signers([values.admin])
+      .rpc({ skipPreflight: true });
+
+    // 2. Second Deposit
+    const secondDepositAmount = new anchor.BN(100000);
+    await program.methods
+      .depositLiquidity(secondDepositAmount, secondDepositAmount)
+      .accounts({
+        pool: values.poolKey,
+        poolAuthority: values.poolAuthority,
+        depositor: values.admin.publicKey,
+        mintLiquidity: values.mintLiquidity,
+        mintA: values.mintAKeypair.publicKey,
+        mintB: values.mintBKeypair.publicKey,
+        poolAccountA: values.poolAccountA,
+        poolAccountB: values.poolAccountB,
+        depositorAccountLiquidity: values.liquidityAccount,
+        depositorAccountA: values.holderAccountA,
+        depositorAccountB: values.holderAccountB,
+      })
+      .signers([values.admin])
+      .rpc({ skipPreflight: true });
+
+    const poolAccountA = await connection.getTokenAccountBalance(values.poolAccountA);
+    expect(poolAccountA.value.amount).to.equal(values.depositAmountA.add(secondDepositAmount).toString());
+  });
+
+  it('Deposit with different ratio', async () => {
+    // 1. Initial Deposit with 1:5 ratio
+    // Pool A: 1,000,000
+    // Pool B: 5,000,000
+    const initialAmountA = new anchor.BN(1_000_000);
+    const initialAmountB = new anchor.BN(5_000_000);
+
+    await program.methods
+      .depositLiquidity(initialAmountA, initialAmountB)
+      .accounts({
+        pool: values.poolKey,
+        poolAuthority: values.poolAuthority,
+        depositor: values.admin.publicKey,
+        mintLiquidity: values.mintLiquidity,
+        mintA: values.mintAKeypair.publicKey,
+        mintB: values.mintBKeypair.publicKey,
+        poolAccountA: values.poolAccountA,
+        poolAccountB: values.poolAccountB,
+        depositorAccountLiquidity: values.liquidityAccount,
+        depositorAccountA: values.holderAccountA,
+        depositorAccountB: values.holderAccountB,
+      })
+      .signers([values.admin])
+      .rpc({ skipPreflight: true });
+
+        // 2. Second Deposit with mismatched input
+        // Input A: 500,000
+        // Input B: 500,000
+        // Logic:
+        // - 500k A requires 2.5M B. (User only provided 500k B).
+        // - 500k B requires 100k A. (User provided 500k A).
+        // Result: Deposit 100k A and 500k B.
+        const secondDepositA = new anchor.BN(500000);
+        const secondDepositBInput = new anchor.BN(500000);
+    
+        await program.methods
+          .depositLiquidity(secondDepositA, secondDepositBInput)
+          .accounts({
+            pool: values.poolKey,
+            poolAuthority: values.poolAuthority,
+            depositor: values.admin.publicKey,
+            mintLiquidity: values.mintLiquidity,
+            mintA: values.mintAKeypair.publicKey,
+            mintB: values.mintBKeypair.publicKey,
+            poolAccountA: values.poolAccountA,
+            poolAccountB: values.poolAccountB,
+            depositorAccountLiquidity: values.liquidityAccount,
+            depositorAccountA: values.holderAccountA,
+            depositorAccountB: values.holderAccountB,
+          })
+          .signers([values.admin])
+          .rpc({ skipPreflight: true });
+    
+        // 3. Verify Balances
+        const poolAccountA = await connection.getTokenAccountBalance(values.poolAccountA);
+        const poolAccountB = await connection.getTokenAccountBalance(values.poolAccountB);
+    
+        // Total A: 1,000,000 + 100,000 = 1,100,000
+        // We expect 100,000 A to be deposited.
+        const expectedAdditionalA = secondDepositBInput.mul(initialAmountA).div(initialAmountB); // 500k * (1M/5M) = 100k
+        expect(poolAccountA.value.amount).to.equal(initialAmountA.add(expectedAdditionalA).toString());
+        
+        // Total B: 5,000,000 + 500,000 = 5,500,000
+        expect(poolAccountB.value.amount).to.equal(initialAmountB.add(secondDepositBInput).toString());
+      });
+    });