stake-pool: Fix stake split and withdraw token calculation (#1352)

* Fix stake split and withdraw token calculation

* Cargo fmt

stake-pool/cli/src/main.rs

@@ -395,7 +395,9 @@ fn command_vsa_remove(
 
     // Calculate amount of tokens to burn
     let stake_account = config.rpc_client.get_account(&stake)?;
-    let tokens_to_burn = stake_amount_to_pool_tokens(&pool_data, stake_account.lamports);
+    let tokens_to_burn = pool_data
+        .calc_pool_withdraw_amount(stake_account.lamports)
+        .unwrap();
 
     // Check balance and mint
     let account_data = config.rpc_client.get_account_data(&burn_from)?;
@@ -736,22 +738,6 @@ fn command_update(config: &Config, pool: &Pubkey) -> CommandResult {
     }
 }
 
-fn stake_amount_to_pool_tokens(pool_data: &StakePool, amount: u64) -> u64 {
-    (amount as u128)
-        .checked_mul(pool_data.pool_total as u128)
-        .unwrap()
-        .checked_div(pool_data.stake_total as u128)
-        .unwrap() as u64
-}
-
-fn pool_tokens_to_stake_amount(pool_data: &StakePool, tokens: u64) -> u64 {
-    (tokens as u128)
-        .checked_mul(pool_data.stake_total as u128)
-        .unwrap()
-        .checked_div(pool_data.pool_total as u128)
-        .unwrap() as u64
-}
-
 #[derive(PartialEq, Debug)]
 struct WithdrawAccount {
     pubkey: Pubkey,
@@ -859,7 +845,7 @@ fn command_withdraw(
     }
 
     // Convert pool tokens amount to lamports
-    let sol_withdraw_amount = pool_tokens_to_stake_amount(&pool_data, amount);
+    let sol_withdraw_amount = pool_data.calc_lamports_amount(amount).unwrap();
 
     // Get the list of accounts to withdraw from
     let withdraw_from: Vec<WithdrawAccount> =
@@ -870,8 +856,6 @@ fn command_withdraw(
     let mut signers = vec![config.fee_payer.as_ref(), config.owner.as_ref()];
     let stake_receiver_account = Keypair::new(); // Will be added to signers if creating new account
 
-    let mut total_rent_free_balances: u64 = 0;
-
     instructions.push(
         // Approve spending token
         approve_token(
@@ -887,12 +871,19 @@ fn command_withdraw(
     // Use separate mutable variable because withdraw might create a new account
     let mut stake_receiver: Option<Pubkey> = *stake_receiver_param;
 
+    let mut total_rent_free_balances = 0;
+
     // Go through prepared accounts and withdraw/claim them
     for withdraw_stake in withdraw_from {
+        let withdraw_amount = pool_data
+            .calc_pool_withdraw_amount(withdraw_stake.amount)
+            .unwrap()
+            + 1;
         println!(
-            "Withdrawing from account {}, amount {} SOL",
+            "Withdrawing from account {}, amount {} SOL, {} pool tokens",
             withdraw_stake.pubkey,
-            lamports_to_sol(withdraw_stake.amount)
+            lamports_to_sol(withdraw_stake.amount),
+            lamports_to_sol(withdraw_amount),
         );
 
         if stake_receiver.is_none() {
@@ -936,7 +927,7 @@ fn command_withdraw(
             &pool_data.pool_mint,
             &spl_token::id(),
             &stake_program_id(),
-            withdraw_stake.amount,
+            withdraw_amount,
         )?);
     }
 

stake-pool/program/src/processor.rs

@@ -131,7 +131,6 @@ impl Processor {
     }
 
     /// Issue a stake_split instruction.
-    #[allow(clippy::too_many_arguments)]
     pub fn stake_split<'a>(
         stake_pool: &Pubkey,
         stake_account: AccountInfo<'a>,
@@ -140,8 +139,6 @@ impl Processor {
         bump_seed: u8,
         amount: u64,
         split_stake: AccountInfo<'a>,
-        reserved: AccountInfo<'a>,
-        stake_program_info: AccountInfo<'a>,
     ) -> Result<(), ProgramError> {
         let me_bytes = stake_pool.to_bytes();
         let authority_signature_seeds = [&me_bytes[..32], authority_type, &[bump_seed]];
@@ -149,17 +146,7 @@ impl Processor {
 
         let ix = stake::split_only(stake_account.key, authority.key, amount, split_stake.key);
 
-        invoke_signed(
-            &ix,
-            &[
-                stake_account,
-                reserved,
-                authority,
-                split_stake,
-                stake_program_info,
-            ],
-            signers,
-        )
+        invoke_signed(&ix, &[stake_account, split_stake, authority], signers)
     }
 
     /// Issue a stake_merge instruction.
@@ -1114,8 +1101,6 @@ impl Processor {
             stake_pool.withdraw_bump_seed,
             stake_amount,
             stake_split_to.clone(),
-            clock_info.clone(),
-            stake_program_info.clone(),
         )?;
 
         Self::stake_authorize(