fix(token-swap): issue LP shares proportionally after fee growth

tokens/token-swap/README.md

@@ -218,9 +218,7 @@ This code ensures that the amounts of tokens A and B being deposited are provide
 
 https://github.com/solana-developers/program-examples/blob/419cb6b6c20e8b1c65711b68a4dde2527725cc1a/tokens/token-swap/anchor/programs/token-swap/src/instructions/deposit_liquidity.rs#L63-L77
 
-This code calculates the amount of liquidity that is about to be deposited into the pool. It calculates the square root of the product of **`amount_a`** and **`amount_b`**, using fixed-point arithmetic to ensure precision.
-
-If this is the first deposit (pool creation), the function checks if the calculated liquidity is greater than the **`MINIMUM_LIQUIDITY`** constant (a minimum liquidity required for the pool). If it's not, the function returns an error to indicate that the deposit is too small. Additionally, it subtracts the **`MINIMUM_LIQUIDITY`** from the calculated liquidity to lock it as the initial liquidity.
+For the first deposit, the program uses the square root of `amount_a * amount_b` and subtracts `MINIMUM_LIQUIDITY`, which remains locked. Later deposits receive shares proportional to the existing reserves and effective LP supply (the minted supply plus the locked minimum). The program rounds down and uses the smaller share calculated from the two deposited amounts, so rounding cannot dilute existing providers. A deposit that would mint no shares is rejected.
 
 https://github.com/solana-developers/program-examples/blob/419cb6b6c20e8b1c65711b68a4dde2527725cc1a/tokens/token-swap/anchor/programs/token-swap/src/instructions/deposit_liquidity.rs#L79-L101
 

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

@@ -10,11 +10,7 @@ use crate::{
     state::Pool,
 };
 
-pub fn deposit_liquidity(
-    ctx: Context<DepositLiquidity>,
-    amount_a: u64,
-    amount_b: u64,
-) -> Result<()> {
+pub fn deposit_liquidity(ctx: Context<DepositLiquidity>, amount_a: u64, amount_b: u64) -> Result<()> {
     // Prevent depositing assets the depositor does not own
     let mut amount_a = if amount_a > ctx.accounts.depositor_account_a.amount {
         ctx.accounts.depositor_account_a.amount
@@ -36,48 +32,47 @@ pub fn deposit_liquidity(
         // Add as is if there is no liquidity
         (amount_a, amount_b)
     } else {
+        require!(pool_a.amount > 0 && pool_b.amount > 0, TutorialError::DepositTooSmall);
         // 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();
+        let amount_b_required = amount_a_u128.checked_mul(pool_b_u128).unwrap().checked_div(pool_a_u128).unwrap();
 
         if amount_b_required <= amount_b_u128 {
             // We have enough B to match the A provided
             (amount_a, amount_b_required as u64)
         } else {
             // 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();
+            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)
         }
     };
 
-    // Computing the amount of liquidity about to be deposited.
-    // Multiply in u128 so the product of two u64 amounts cannot overflow.
-    let mut liquidity = (amount_a as u128)
-        .checked_mul(amount_b as u128)
-        .unwrap()
-        .isqrt() as u64;
-
-    // Lock some minimum liquidity on the first deposit
-    if pool_creation {
-        if liquidity < MINIMUM_LIQUIDITY {
-            return err!(TutorialError::DepositTooSmall);
-        }
-
-        liquidity -= MINIMUM_LIQUIDITY;
-    }
+    let liquidity = if pool_creation {
+        // The initial share supply is the geometric mean, less the locked minimum.
+        let initial = ((amount_a as u128) * (amount_b as u128)).isqrt() as u64;
+        initial.checked_sub(MINIMUM_LIQUIDITY).ok_or(TutorialError::DepositTooSmall)?
+    } else {
+        // Withdrawal accounts for the unminted locked minimum as part of supply.
+        // Limit shares by both actual deposits so rounding cannot dilute existing LPs.
+        let supply = ctx.accounts.mint_liquidity.supply as u128 + MINIMUM_LIQUIDITY as u128;
+        let shares_a = (amount_a as u128)
+            .checked_mul(supply)
+            .ok_or(TutorialError::MathOverflow)?
+            .checked_div(pool_a.amount as u128)
+            .ok_or(TutorialError::MathOverflow)?;
+        let shares_b = (amount_b as u128)
+            .checked_mul(supply)
+            .ok_or(TutorialError::MathOverflow)?
+            .checked_div(pool_b.amount as u128)
+            .ok_or(TutorialError::MathOverflow)?;
+        u64::try_from(shares_a.min(shares_b)).map_err(|_| TutorialError::MathOverflow)?
+    };
+    require!(liquidity > 0, TutorialError::DepositTooSmall);
 
     // Transfer tokens to the pool
     token::transfer(

tokens/token-swap/anchor/programs/token-swap/tests/test_token_swap.rs

@@ -159,30 +159,35 @@ impl Env {
     }
 
     fn deposit(&mut self, amount_a: u64, amount_b: u64) {
+        let admin = self.admin.insecure_clone();
+        self.deposit_as(&admin, amount_a, amount_b);
+    }
+
+    fn deposit_as(&mut self, depositor: &Keypair, amount_a: u64, amount_b: u64) {
         let keys = self.keys();
         let ix = Instruction::new_with_bytes(
             swap_example::id(),
             &swap_example::instruction::DepositLiquidity { amount_a, amount_b }.data(),
             swap_example::accounts::DepositLiquidity {
                 pool: keys.pool,
                 pool_authority: keys.pool_authority,
-                depositor: self.admin.pubkey(),
+                depositor: depositor.pubkey(),
                 mint_liquidity: keys.mint_liquidity,
                 mint_a: self.mint_a.pubkey(),
                 mint_b: self.mint_b.pubkey(),
                 pool_account_a: keys.pool_account_a,
                 pool_account_b: keys.pool_account_b,
-                depositor_account_liquidity: self.liquidity_account(),
-                depositor_account_a: self.holder_a(),
-                depositor_account_b: self.holder_b(),
+                depositor_account_liquidity: get_associated_token_address(&depositor.pubkey(), &keys.mint_liquidity),
+                depositor_account_a: get_associated_token_address(&depositor.pubkey(), &self.mint_a.pubkey()),
+                depositor_account_b: get_associated_token_address(&depositor.pubkey(), &self.mint_b.pubkey()),
                 payer: self.payer.pubkey(),
                 token_program: spl_token::ID,
                 associated_token_program: associated_token::ID,
                 system_program: system_program::ID,
             }
             .to_account_metas(None),
         );
-        send(&mut self.svm, &self.payer, &[ix], &[&self.admin]).unwrap();
+        send(&mut self.svm, &self.payer, &[ix], &[depositor]).unwrap();
     }
 
     fn withdraw(&mut self, amount: u64) {
@@ -354,6 +359,48 @@ fn deposit_with_different_ratio_is_limited_by_pool_ratio() {
     assert_eq!(env.balance(&keys.pool_account_b), initial_b + second_b);
 }
 
+#[test]
+fn deposit_after_fee_growth_preserves_existing_lp_share() {
+    let mut env = Env::with_pool();
+    env.deposit(4_000_000, 4_000_000);
+    env.swap(true, 2_000_000, 1);
+
+    let keys = env.keys();
+    let before_a = env.balance(&keys.pool_account_a);
+    let before_b = env.balance(&keys.pool_account_b);
+    let existing_lp = env.balance(&env.liquidity_account());
+    let new_lp = Keypair::new();
+    let new_lp_account = get_associated_token_address(&new_lp.pubkey(), &keys.mint_liquidity);
+    let mut setup = Vec::new();
+    for (mint, source) in [(env.mint_a.pubkey(), env.holder_a()), (env.mint_b.pubkey(), env.holder_b())] {
+        let destination = get_associated_token_address(&new_lp.pubkey(), &mint);
+        setup.push(create_associated_token_account(&env.payer.pubkey(), &new_lp.pubkey(), &mint, &spl_token::ID));
+        setup.push(
+            spl_token::instruction::transfer(
+                &spl_token::ID,
+                &source,
+                &destination,
+                &env.admin.pubkey(),
+                &[],
+                2_000_000,
+            )
+            .unwrap(),
+        );
+    }
+    send(&mut env.svm, &env.payer, &setup, &[&env.admin]).unwrap();
+    env.deposit_as(&new_lp, 1_000_000, 1_000_000);
+
+    let added_a = env.balance(&keys.pool_account_a) - before_a;
+    let added_b = env.balance(&keys.pool_account_b) - before_b;
+    let effective_supply = existing_lp + MINIMUM_LIQUIDITY;
+    let max_lp_a = (added_a as u128 * effective_supply as u128 / before_a as u128) as u64;
+    let max_lp_b = (added_b as u128 * effective_supply as u128 / before_b as u128) as u64;
+    let minted_lp = env.balance(&new_lp_account);
+    assert!(minted_lp > 0);
+    assert_eq!(minted_lp, max_lp_a.min(max_lp_b), "new shares must reflect the actual contribution to both reserves");
+    assert_eq!(env.balance(&env.liquidity_account()), existing_lp);
+}
+
 #[test]
 fn swap_a_to_b() {
     let mut env = Env::with_pool();