stake-pool: Use checked_ceil_div for withdraw calc (#1482)

* stake-pool: Use checked_ceil_div for withdraw calc

When a stake account is totally removed from a stake pool by the
manager, there's a chance that the operation would not take enough of
the manager's pool tokens by 1 due to truncation.

Do a ceiling division instead, and refactor ceiling division into the
math library.

* Use new function name on CLI

* Cargo fmt

Cargo.lock

@@ -3865,6 +3865,7 @@ dependencies = [
  "solana-program-test",
  "solana-sdk",
  "solana-vote-program",
+ "spl-math",
  "spl-token 3.1.0",
  "thiserror",
 ]

libraries/math/src/checked_ceil_div.rs

@@ -1,6 +1,6 @@
-//! Defines useful math utils
+//! Defines performing checked ceiling division for different types
 
-use spl_math::uint::U256;
+use crate::uint::U256;
 
 /// Perform a division that does not truncate value from either side, returning
 /// the (quotient, divisor) as a tuple

libraries/math/src/lib.rs

@@ -4,6 +4,7 @@
 #![forbid(unsafe_code)]
 
 pub mod approximations;
+pub mod checked_ceil_div;
 mod entrypoint;
 pub mod error;
 pub mod instruction;

stake-pool/cli/src/main.rs

@@ -777,7 +777,7 @@ fn pick_withdraw_accounts(
             continue;
         }
         let available_for_withdrawal = stake_pool
-            .calc_lamports_amount(account.lamports - *MIN_STAKE_BALANCE)
+            .calc_lamports_withdraw_amount(account.lamports - *MIN_STAKE_BALANCE)
             .unwrap();
         let withdraw_amount = u64::min(available_for_withdrawal, remaining_amount);
 
@@ -878,7 +878,7 @@ fn command_withdraw(
     for withdraw_stake in withdraw_accounts {
         // Convert pool tokens amount to lamports
         let sol_withdraw_amount = pool_data
-            .calc_lamports_amount(withdraw_stake.pool_amount)
+            .calc_lamports_withdraw_amount(withdraw_stake.pool_amount)
             .unwrap();
 
         println!(

stake-pool/program/Cargo.toml

@@ -19,6 +19,7 @@ num_enum = "0.5.1"
 serde = "1.0.121"
 serde_derive = "1.0.103"
 solana-program = "1.6.1"
+spl-math = { path = "../../libraries/math", features = [ "no-entrypoint" ] }
 spl-token = { path = "../../token/program", features = [ "no-entrypoint" ] }
 thiserror = "1.0"
 bincode = "1.3.1"

stake-pool/program/src/processor.rs

@@ -1090,7 +1090,7 @@ impl Processor {
             .ok_or(StakePoolError::ValidatorNotFound)?;
 
         let stake_amount = stake_pool
-            .calc_lamports_amount(pool_amount)
+            .calc_lamports_withdraw_amount(pool_amount)
             .ok_or(StakePoolError::CalculationFailure)?;
 
         Self::stake_split(

stake-pool/program/src/state.rs

@@ -6,6 +6,7 @@ use {
         account_info::AccountInfo, entrypoint::ProgramResult, program_error::ProgramError,
         pubkey::Pubkey,
     },
+    spl_math::checked_ceil_div::CheckedCeilDiv,
     std::convert::{TryFrom, TryInto},
     std::mem::size_of,
 };
@@ -50,10 +51,6 @@ impl StakePool {
         if self.stake_total == 0 {
             return Some(stake_lamports);
         }
-        self.calc_pool_withdraw_amount(stake_lamports)
-    }
-    /// calculate the pool tokens that should be withdrawn
-    pub fn calc_pool_withdraw_amount(&self, stake_lamports: u64) -> Option<u64> {
         u64::try_from(
             (stake_lamports as u128)
                 .checked_mul(self.pool_total as u128)?
@@ -61,8 +58,15 @@ impl StakePool {
         )
         .ok()
     }
-    /// calculate lamports amount
-    pub fn calc_lamports_amount(&self, pool_tokens: u64) -> Option<u64> {
+    /// calculate the pool tokens that should be withdrawn
+    pub fn calc_pool_withdraw_amount(&self, stake_lamports: u64) -> Option<u64> {
+        let (quotient, _) = (stake_lamports as u128)
+            .checked_mul(self.pool_total as u128)?
+            .checked_ceil_div(self.stake_total as u128)?;
+        u64::try_from(quotient).ok()
+    }
+    /// calculate lamports amount on withdrawal
+    pub fn calc_lamports_withdraw_amount(&self, pool_tokens: u64) -> Option<u64> {
         u64::try_from(
             (pool_tokens as u128)
                 .checked_mul(self.stake_total as u128)?

token-swap/program/src/curve/constant_price.rs

@@ -1,19 +1,20 @@
 //! Simple constant price swap curve, set at init
 
-use crate::{
-    curve::calculator::{
-        map_zero_to_none, CurveCalculator, DynPack, RoundDirection, SwapWithoutFeesResult,
-        TradeDirection, TradingTokenResult,
+use {
+    crate::{
+        curve::calculator::{
+            map_zero_to_none, CurveCalculator, DynPack, RoundDirection, SwapWithoutFeesResult,
+            TradeDirection, TradingTokenResult,
+        },
+        error::SwapError,
     },
-    curve::math::CheckedCeilDiv,
-    error::SwapError,
+    arrayref::{array_mut_ref, array_ref},
+    solana_program::{
+        program_error::ProgramError,
+        program_pack::{IsInitialized, Pack, Sealed},
+    },
+    spl_math::{checked_ceil_div::CheckedCeilDiv, precise_number::PreciseNumber, uint::U256},
 };
-use arrayref::{array_mut_ref, array_ref};
-use solana_program::{
-    program_error::ProgramError,
-    program_pack::{IsInitialized, Pack, Sealed},
-};
-use spl_math::{precise_number::PreciseNumber, uint::U256};
 
 /// ConstantPriceCurve struct implementing CurveCalculator
 #[derive(Clone, Debug, Default, PartialEq)]

token-swap/program/src/curve/constant_product.rs

@@ -6,14 +6,13 @@ use {
             map_zero_to_none, CurveCalculator, DynPack, RoundDirection, SwapWithoutFeesResult,
             TradeDirection, TradingTokenResult,
         },
-        curve::math::CheckedCeilDiv,
         error::SwapError,
     },
     solana_program::{
         program_error::ProgramError,
         program_pack::{IsInitialized, Pack, Sealed},
     },
-    spl_math::precise_number::PreciseNumber,
+    spl_math::{checked_ceil_div::CheckedCeilDiv, precise_number::PreciseNumber},
 };
 
 /// ConstantProductCurve struct implementing CurveCalculator

token-swap/program/src/curve/mod.rs

@@ -5,6 +5,5 @@ pub mod calculator;
 pub mod constant_price;
 pub mod constant_product;
 pub mod fees;
-pub mod math;
 pub mod offset;
 pub mod stable;