stake-pool-js: Remove divideBnToNumber, use BNs more (#6482)

* stake-pool-js: Remove `divideBnToNumber`, use BNs more

* Add defensive check for negative number

* Revert ceil on calc

Author: joncinque

stake-pool/js/src/index.ts

@@ -60,7 +60,7 @@ export interface StakePoolAccount {
 export interface WithdrawAccount {
   stakeAddress: PublicKey;
   voteAddress?: PublicKey;
-  poolAmount: number;
+  poolAmount: BN;
 }
 
 /**
@@ -354,7 +354,7 @@ export async function withdrawStake(
   validatorComparator?: (_a: ValidatorAccount, _b: ValidatorAccount) => number,
 ) {
   const stakePool = await getStakePoolAccount(connection, stakePoolAddress);
-  const poolAmount = solToLamports(amount);
+  const poolAmount = new BN(solToLamports(amount));
 
   if (!poolTokenAccount) {
     poolTokenAccount = getAssociatedTokenAddressSync(stakePool.account.data.poolMint, tokenOwner);
@@ -363,7 +363,7 @@ export async function withdrawStake(
   const tokenAccount = await getAccount(connection, poolTokenAccount);
 
   // Check withdrawFrom balance
-  if (tokenAccount.amount < poolAmount) {
+  if (tokenAccount.amount < poolAmount.toNumber()) {
     throw new Error(
       `Not enough token balance to withdraw ${lamportsToSol(poolAmount)} pool tokens.
         Maximum withdraw amount is ${lamportsToSol(tokenAccount.amount)} pool tokens.`,
@@ -420,10 +420,10 @@ export async function withdrawStake(
 
       const availableForWithdrawal = calcLamportsWithdrawAmount(
         stakePool.account.data,
-        stakeAccount.lamports - MINIMUM_ACTIVE_STAKE - stakeAccountRentExemption,
+        new BN(stakeAccount.lamports - MINIMUM_ACTIVE_STAKE - stakeAccountRentExemption),
       );
 
-      if (availableForWithdrawal < poolAmount) {
+      if (availableForWithdrawal.lt(poolAmount)) {
         throw new Error(
           `Not enough lamports available for withdrawal from ${stakeAccountAddress},
             ${poolAmount} asked, ${availableForWithdrawal} available.`,
@@ -450,12 +450,18 @@ export async function withdrawStake(
       throw new Error('Invalid Stake Account');
     }
 
+    const availableLamports = new BN(
+      stakeAccount.lamports - MINIMUM_ACTIVE_STAKE - stakeAccountRentExemption,
+    );
+    if (availableLamports.lt(new BN(0))) {
+      throw new Error('Invalid Stake Account');
+    }
     const availableForWithdrawal = calcLamportsWithdrawAmount(
       stakePool.account.data,
-      stakeAccount.lamports - MINIMUM_ACTIVE_STAKE - stakeAccountRentExemption,
+      availableLamports,
     );
 
-    if (availableForWithdrawal < poolAmount) {
+    if (availableForWithdrawal.lt(poolAmount)) {
       // noinspection ExceptionCaughtLocallyJS
       throw new Error(
         `Not enough lamports available for withdrawal from ${stakeAccountAddress},
@@ -492,7 +498,7 @@ export async function withdrawStake(
       poolTokenAccount,
       userTransferAuthority.publicKey,
       tokenOwner,
-      poolAmount,
+      poolAmount.toNumber(),
     ),
   );
 
@@ -508,8 +514,9 @@ export async function withdrawStake(
       break;
     }
     // Convert pool tokens amount to lamports
-    const solWithdrawAmount = Math.ceil(
-      calcLamportsWithdrawAmount(stakePool.account.data, withdrawAccount.poolAmount),
+    const solWithdrawAmount = calcLamportsWithdrawAmount(
+      stakePool.account.data,
+      withdrawAccount.poolAmount,
     );
 
     let infoMsg = `Withdrawing ◎${solWithdrawAmount},
@@ -542,7 +549,7 @@ export async function withdrawStake(
         sourcePoolAccount: poolTokenAccount,
         managerFeeAccount: stakePool.account.data.managerFeeAccount,
         poolMint: stakePool.account.data.poolMint,
-        poolTokens: withdrawAccount.poolAmount,
+        poolTokens: withdrawAccount.poolAmount.toNumber(),
         withdrawAuthority,
       }),
     );

stake-pool/js/src/layouts.ts

@@ -1,5 +1,5 @@
 import { publicKey, struct, u32, u64, u8, option, vec } from '@coral-xyz/borsh';
-import { Lockup, PublicKey } from '@solana/web3.js';
+import { PublicKey } from '@solana/web3.js';
 import BN from 'bn.js';
 import {
   Infer,
@@ -76,6 +76,11 @@ export const StakeAccount = type({
   type: StakeAccountType,
   info: optional(StakeAccountInfo),
 });
+export interface Lockup {
+  unixTimestamp: BN;
+  epoch: BN;
+  custodian: PublicKey;
+}
 
 export interface StakePool {
   accountType: AccountType;

stake-pool/js/src/utils/stake.ts

@@ -41,14 +41,14 @@ export interface ValidatorAccount {
   type: 'preferred' | 'active' | 'transient' | 'reserve';
   voteAddress?: PublicKey | undefined;
   stakeAddress: PublicKey;
-  lamports: number;
+  lamports: BN;
 }
 
 export async function prepareWithdrawAccounts(
   connection: Connection,
   stakePool: StakePool,
   stakePoolAddress: PublicKey,
-  amount: number,
+  amount: BN,
   compareFn?: (a: ValidatorAccount, b: ValidatorAccount) => number,
   skipFee?: boolean,
 ): Promise<WithdrawAccount[]> {
@@ -62,13 +62,13 @@ export async function prepareWithdrawAccounts(
   const minBalanceForRentExemption = await connection.getMinimumBalanceForRentExemption(
     StakeProgram.space,
   );
-  const minBalance = minBalanceForRentExemption + MINIMUM_ACTIVE_STAKE;
+  const minBalance = new BN(minBalanceForRentExemption + MINIMUM_ACTIVE_STAKE);
 
   let accounts = [] as Array<{
     type: 'preferred' | 'active' | 'transient' | 'reserve';
     voteAddress?: PublicKey | undefined;
     stakeAddress: PublicKey;
-    lamports: number;
+    lamports: BN;
   }>;
 
   // Prepare accounts
@@ -91,12 +91,12 @@ export async function prepareWithdrawAccounts(
         type: isPreferred ? 'preferred' : 'active',
         voteAddress: validator.voteAccountAddress,
         stakeAddress: stakeAccountAddress,
-        lamports: validator.activeStakeLamports.toNumber(),
+        lamports: validator.activeStakeLamports,
       });
     }
 
-    const transientStakeLamports = validator.transientStakeLamports.toNumber() - minBalance;
-    if (transientStakeLamports > 0) {
+    const transientStakeLamports = validator.transientStakeLamports.sub(minBalance);
+    if (transientStakeLamports.gt(new BN(0))) {
       const transientStakeAccountAddress = await findTransientStakeProgramAddress(
         STAKE_POOL_PROGRAM_ID,
         validator.voteAccountAddress,
@@ -113,11 +113,11 @@ export async function prepareWithdrawAccounts(
   }
 
   // Sort from highest to lowest balance
-  accounts = accounts.sort(compareFn ? compareFn : (a, b) => b.lamports - a.lamports);
+  accounts = accounts.sort(compareFn ? compareFn : (a, b) => b.lamports.sub(a.lamports).toNumber());
 
   const reserveStake = await connection.getAccountInfo(stakePool.reserveStake);
-  const reserveStakeBalance = (reserveStake?.lamports ?? 0) - minBalanceForRentExemption;
-  if (reserveStakeBalance > 0) {
+  const reserveStakeBalance = new BN((reserveStake?.lamports ?? 0) - minBalanceForRentExemption);
+  if (reserveStakeBalance.gt(new BN(0))) {
     accounts.push({
       type: 'reserve',
       stakeAddress: stakePool.reserveStake,
@@ -127,7 +127,7 @@ export async function prepareWithdrawAccounts(
 
   // Prepare the list of accounts to withdraw from
   const withdrawFrom: WithdrawAccount[] = [];
-  let remainingAmount = amount;
+  let remainingAmount = new BN(amount);
 
   const fee = stakePool.stakeWithdrawalFee;
   const inverseFee: Fee = {
@@ -139,40 +139,39 @@ export async function prepareWithdrawAccounts(
     const filteredAccounts = accounts.filter((a) => a.type == type);
 
     for (const { stakeAddress, voteAddress, lamports } of filteredAccounts) {
-      if (lamports <= minBalance && type == 'transient') {
+      if (lamports.lte(minBalance) && type == 'transient') {
         continue;
       }
 
       let availableForWithdrawal = calcPoolTokensForDeposit(stakePool, lamports);
 
       if (!skipFee && !inverseFee.numerator.isZero()) {
-        availableForWithdrawal = divideBnToNumber(
-          new BN(availableForWithdrawal).mul(inverseFee.denominator),
-          inverseFee.numerator,
-        );
+        availableForWithdrawal = availableForWithdrawal
+          .mul(inverseFee.denominator)
+          .div(inverseFee.numerator);
       }
 
-      const poolAmount = Math.min(availableForWithdrawal, remainingAmount);
-      if (poolAmount <= 0) {
+      const poolAmount = BN.min(availableForWithdrawal, remainingAmount);
+      if (poolAmount.lte(new BN(0))) {
         continue;
       }
 
       // Those accounts will be withdrawn completely with `claim` instruction
       withdrawFrom.push({ stakeAddress, voteAddress, poolAmount });
-      remainingAmount -= poolAmount;
+      remainingAmount = remainingAmount.sub(poolAmount);
 
-      if (remainingAmount == 0) {
+      if (remainingAmount.isZero()) {
         break;
       }
     }
 
-    if (remainingAmount == 0) {
+    if (remainingAmount.isZero()) {
       break;
     }
   }
 
   // Not enough stake to withdraw the specified amount
-  if (remainingAmount > 0) {
+  if (remainingAmount.gt(new BN(0))) {
     throw new Error(
       `No stake accounts found in this pool with enough balance to withdraw ${lamportsToSol(
         amount,
@@ -186,35 +185,24 @@ export async function prepareWithdrawAccounts(
 /**
  * Calculate the pool tokens that should be minted for a deposit of `stakeLamports`
  */
-export function calcPoolTokensForDeposit(stakePool: StakePool, stakeLamports: number): number {
+export function calcPoolTokensForDeposit(stakePool: StakePool, stakeLamports: BN): BN {
   if (stakePool.poolTokenSupply.isZero() || stakePool.totalLamports.isZero()) {
     return stakeLamports;
   }
-  return Math.floor(
-    divideBnToNumber(new BN(stakeLamports).mul(stakePool.poolTokenSupply), stakePool.totalLamports),
-  );
+  const numerator = stakeLamports.mul(stakePool.poolTokenSupply);
+  return numerator.div(stakePool.totalLamports);
 }
 
 /**
  * Calculate lamports amount on withdrawal
  */
-export function calcLamportsWithdrawAmount(stakePool: StakePool, poolTokens: number): number {
-  const numerator = new BN(poolTokens).mul(stakePool.totalLamports);
+export function calcLamportsWithdrawAmount(stakePool: StakePool, poolTokens: BN): BN {
+  const numerator = poolTokens.mul(stakePool.totalLamports);
   const denominator = stakePool.poolTokenSupply;
   if (numerator.lt(denominator)) {
-    return 0;
-  }
-  return divideBnToNumber(numerator, denominator);
-}
-
-export function divideBnToNumber(numerator: BN, denominator: BN): number {
-  if (denominator.isZero()) {
-    return 0;
+    return new BN(0);
   }
-  const quotient = numerator.div(denominator).toNumber();
-  const rem = numerator.umod(denominator);
-  const gcd = rem.gcd(denominator);
-  return quotient + rem.div(gcd).toNumber() / denominator.div(gcd).toNumber();
+  return numerator.div(denominator);
 }
 
 export function newStakeAccount(

stake-pool/js/test/calculation.test.ts

@@ -0,0 +1,15 @@
+import { LAMPORTS_PER_SOL } from '@solana/web3.js';
+import { stakePoolMock } from './mocks';
+import { calcPoolTokensForDeposit } from '../src/utils/stake';
+import BN from 'bn.js';
+
+describe('calculations', () => {
+  it('should successfully calculate pool tokens for a pool with a lot of stake', () => {
+    const lamports = new BN(LAMPORTS_PER_SOL * 100);
+    const bigStakePoolMock = stakePoolMock;
+    bigStakePoolMock.totalLamports = new BN('11000000000000000'); // 11 million SOL
+    bigStakePoolMock.poolTokenSupply = new BN('10000000000000000'); // 10 million tokens
+    const availableForWithdrawal = calcPoolTokensForDeposit(bigStakePoolMock, lamports);
+    expect(availableForWithdrawal.toNumber()).toEqual(90909090909);
+  });
+});