Merge pull request #486 from balancer/bpt-from-reference-amount-boosted

Add proportional helper to calculate BPT from a reference amount for boosted pools

.changeset/nice-balloons-train.md

@@ -0,0 +1,5 @@
+---
+"@balancer/sdk": patch
+---
+
+Add proportional helper to calculate BPT from a reference amount for boosted pools

src/entities/addLiquidityBoosted/index.ts

@@ -11,7 +11,11 @@ import { getAmountsCall } from '../addLiquidity/helpers';
 
 import { PoolStateWithUnderlyings } from '@/entities/types';
 
-import { getAmounts, getSortedTokens } from '@/entities/utils';
+import {
+    getAmounts,
+    getBptAmountFromReferenceAmountBoosted,
+    getSortedTokens,
+} from '@/entities/utils';
 
 import {
     AddLiquidityBuildCallOutput,
@@ -89,10 +93,10 @@ export class AddLiquidityBoostedV3 {
                 break;
             }
             case AddLiquidityKind.Proportional: {
-                if (input.referenceAmount.address !== poolState.address) {
-                    // TODO: add getBptAmountFromReferenceAmount
-                    throw new Error('Reference token must be the pool token');
-                }
+                const bptAmount = await getBptAmountFromReferenceAmountBoosted(
+                    input,
+                    poolState,
+                );
 
                 const exactAmountsInNumbers =
                     await doAddLiquidityProportionalQuery(
@@ -101,7 +105,7 @@ export class AddLiquidityBoostedV3 {
                         input.sender ?? zeroAddress,
                         input.userData ?? '0x',
                         poolState.address,
-                        input.referenceAmount.rawAmount,
+                        bptAmount.rawAmount,
                     );
 
                 // Amounts are mapped to child tokens of the pool
@@ -114,7 +118,7 @@ export class AddLiquidityBoostedV3 {
 
                 bptOut = TokenAmount.fromRawAmount(
                     bptToken,
-                    input.referenceAmount.rawAmount,
+                    bptAmount.rawAmount,
                 );
                 break;
             }

src/entities/types.ts

@@ -23,6 +23,10 @@ export type PoolTokenWithUnderlying = MinimalToken & {
     underlyingToken: MinimalToken | null;
 };
 
+export interface PoolTokenWithUnderlyingBalance extends PoolTokenWithBalance {
+    underlyingToken: PoolTokenWithBalance | null;
+}
+
 export type PoolStateWithUnderlyings = {
     id: Hex;
     address: Address;
@@ -31,6 +35,15 @@ export type PoolStateWithUnderlyings = {
     protocolVersion: 1 | 2 | 3;
 };
 
+export type PoolStateWithUnderlyingBalances = {
+    id: Hex;
+    address: Address;
+    type: string;
+    tokens: PoolTokenWithUnderlyingBalance[];
+    totalShares: HumanAmount;
+    protocolVersion: 1 | 2 | 3;
+};
+
 export type AddLiquidityAmounts = {
     maxAmountsIn: bigint[];
     maxAmountsInWithoutBpt: bigint[];

src/entities/utils/getPoolStateWithBalancesV3.ts

@@ -1,10 +1,23 @@
-import { createPublicClient, formatEther, formatUnits, http } from 'viem';
+import {
+    createPublicClient,
+    erc4626Abi,
+    formatEther,
+    formatUnits,
+    http,
+    PublicClient,
+} from 'viem';
 
 import { HumanAmount } from '@/data';
 import { CHAINS, VAULT_V3 } from '@/utils';
 
 import { getSortedTokens } from './getSortedTokens';
-import { PoolState, PoolStateWithBalances } from '../types';
+import {
+    PoolState,
+    PoolStateWithBalances,
+    PoolStateWithUnderlyingBalances,
+    PoolStateWithUnderlyings,
+    PoolTokenWithUnderlying,
+} from '../types';
 import { vaultExtensionAbi_V3 } from '@/abi';
 import { TokenAmount } from '../tokenAmount';
 
@@ -13,6 +26,131 @@ export const getPoolStateWithBalancesV3 = async (
     chainId: number,
     rpcUrl: string,
 ): Promise<PoolStateWithBalances> => {
+    const publicClient = createPublicClient({
+        transport: http(rpcUrl),
+        chain: CHAINS[chainId],
+    });
+
+    // get on-chain pool token balances and total shares
+    const { tokenAmounts, totalShares } = await getTokenAmountsAndTotalShares(
+        chainId,
+        poolState,
+        publicClient,
+    );
+
+    // build PoolStateWithBalances object with queried on-chain balances
+    const poolStateWithBalances: PoolStateWithBalances = {
+        ...poolState,
+        tokens: tokenAmounts.map((tokenAmount, i) => ({
+            address: tokenAmount.token.address,
+            decimals: tokenAmount.token.decimals,
+            index: i,
+            balance: formatUnits(
+                tokenAmount.amount,
+                tokenAmount.token.decimals,
+            ) as HumanAmount,
+        })),
+        totalShares: formatEther(totalShares) as HumanAmount,
+    };
+    return poolStateWithBalances;
+};
+
+export const getBoostedPoolStateWithBalancesV3 = async (
+    poolState: PoolStateWithUnderlyings,
+    chainId: number,
+    rpcUrl: string,
+): Promise<PoolStateWithUnderlyingBalances> => {
+    const publicClient = createPublicClient({
+        transport: http(rpcUrl),
+        chain: CHAINS[chainId],
+    });
+
+    // get on-chain pool token balances and total shares
+    const { tokenAmounts, totalShares } = await getTokenAmountsAndTotalShares(
+        chainId,
+        poolState,
+        publicClient,
+    );
+
+    const sortedTokens = [...poolState.tokens].sort(
+        (a, b) => a.index - b.index,
+    );
+
+    // get on-chain balances for each respective underlying pool token (by querying erc4626 previewRedeem function)
+    const underlyingBalances = await getUnderlyingBalances(
+        sortedTokens,
+        tokenAmounts,
+        publicClient,
+    );
+
+    // build PoolStateWithUnderlyingBalances object with queried on-chain balances
+    const poolStateWithUnderlyingBalances: PoolStateWithUnderlyingBalances = {
+        ...poolState,
+        tokens: sortedTokens.map((token, i) => ({
+            ...token,
+            balance: formatUnits(
+                tokenAmounts[i].amount,
+                token.decimals,
+            ) as HumanAmount,
+            underlyingToken:
+                token.underlyingToken === null
+                    ? null
+                    : {
+                          ...token.underlyingToken,
+                          balance: formatUnits(
+                              underlyingBalances.shift() as bigint,
+                              token.underlyingToken.decimals,
+                          ) as HumanAmount,
+                      },
+        })),
+        totalShares: formatEther(totalShares) as HumanAmount,
+    };
+    return poolStateWithUnderlyingBalances;
+};
+
+const getUnderlyingBalances = async (
+    sortedTokens: PoolTokenWithUnderlying[],
+    tokenAmounts: TokenAmount[],
+    publicClient: PublicClient,
+) => {
+    // create one contract call for each underlying token
+    const getUnderlyingBalancesContracts = sortedTokens
+        .filter((token) => token.underlyingToken !== null)
+        .map((token, i) => ({
+            address: token.address,
+            abi: erc4626Abi,
+            functionName: 'previewRedeem' as const,
+            args: [tokenAmounts[i].amount] as const,
+        }));
+
+    // execute multicall to get on-chain balances for each underlying token
+    const underlyingBalanceOutputs = await publicClient.multicall({
+        contracts: [...getUnderlyingBalancesContracts],
+    });
+
+    // throw error if any of the underlying balance calls failed
+    if (
+        underlyingBalanceOutputs.some((output) => output.status === 'failure')
+    ) {
+        throw new Error(
+            'Error: Unable to get underlying balances for v3 pool.',
+        );
+    }
+
+    // extract underlying balances from multicall outputs
+    const underlyingBalances = underlyingBalanceOutputs.map(
+        (output) => output.result as bigint,
+    );
+
+    return underlyingBalances;
+};
+
+const getTokenAmountsAndTotalShares = async (
+    chainId: number,
+    poolState: PoolState,
+    publicClient: PublicClient,
+) => {
+    // create contract calls to get total supply and balances for each pool token
     const totalSupplyContract = {
         address: VAULT_V3[chainId],
         abi: vaultExtensionAbi_V3,
@@ -26,40 +164,25 @@ export const getPoolStateWithBalancesV3 = async (
         args: [poolState.address] as const,
     };
 
-    const publicClient = createPublicClient({
-        transport: http(rpcUrl),
-        chain: CHAINS[chainId],
-    });
+    // execute multicall to get total supply and balances for each pool token
     const outputs = await publicClient.multicall({
         contracts: [totalSupplyContract, getBalanceContracts],
     });
 
+    // throw error if any of the calls failed
     if (outputs.some((output) => output.status === 'failure')) {
         throw new Error(
             'Error: Unable to get pool state with balances for v3 pool.',
         );
     }
 
+    // extract total supply and balances from multicall outputs
     const totalShares = outputs[0].result as bigint;
     const balancesScale18 = outputs[1].result as bigint[];
-
-    const sortedTokens = getSortedTokens(poolState.tokens, chainId);
-    const balances = sortedTokens.map((token, i) =>
+    const poolTokens = getSortedTokens(poolState.tokens, chainId);
+    const tokenAmounts = poolTokens.map((token, i) =>
         TokenAmount.fromScale18Amount(token, balancesScale18[i]),
     );
 
-    const poolStateWithBalances: PoolStateWithBalances = {
-        ...poolState,
-        tokens: sortedTokens.map((token, i) => ({
-            address: token.address,
-            decimals: token.decimals,
-            index: i,
-            balance: formatUnits(
-                balances[i].amount,
-                token.decimals,
-            ) as HumanAmount,
-        })),
-        totalShares: formatEther(totalShares) as HumanAmount,
-    };
-    return poolStateWithBalances;
+    return { tokenAmounts, totalShares };
 };

src/entities/utils/proportionalAmountsHelpers.ts

@@ -1,11 +1,15 @@
 import { Address, parseUnits } from 'viem';
 import { InputAmount } from '@/types';
 import { HumanAmount } from '@/data';
-import { MathSol } from '@/utils';
+import { isSameAddress, MathSol } from '@/utils';
 import { AddLiquidityProportionalInput } from '../addLiquidity/types';
-import { PoolState } from '../types';
+import { PoolState, PoolStateWithUnderlyings } from '../types';
 import { getPoolStateWithBalancesV2 } from './getPoolStateWithBalancesV2';
-import { getPoolStateWithBalancesV3 } from './getPoolStateWithBalancesV3';
+import {
+    getBoostedPoolStateWithBalancesV3,
+    getPoolStateWithBalancesV3,
+} from './getPoolStateWithBalancesV3';
+import { AddLiquidityBoostedProportionalInput } from '../addLiquidityBoosted/types';
 
 /**
  * For a given pool and reference token amount, calculate all token amounts proportional to their balances within the pool.
@@ -127,3 +131,46 @@ export const getBptAmountFromReferenceAmount = async (
     }
     return bptAmount;
 };
+
+/**
+ * Calculate the BPT amount for a given reference amount in a boosted pool (rounded down).
+ *
+ * @param input
+ * @param poolState
+ * @returns
+ */
+export const getBptAmountFromReferenceAmountBoosted = async (
+    input: AddLiquidityBoostedProportionalInput,
+    poolStateWithUnderlyings: PoolStateWithUnderlyings,
+): Promise<InputAmount> => {
+    let bptAmount: InputAmount;
+    if (
+        isSameAddress(
+            input.referenceAmount.address,
+            poolStateWithUnderlyings.address,
+        )
+    ) {
+        bptAmount = input.referenceAmount;
+    } else {
+        const poolStateWithUnderlyingBalances =
+            await getBoostedPoolStateWithBalancesV3(
+                poolStateWithUnderlyings,
+                input.chainId,
+                input.rpcUrl,
+            );
+
+        // use underlying tokens as tokens if they exist (in case of a partial boosted pool)
+        const poolStateWithBalances = {
+            ...poolStateWithUnderlyingBalances,
+            tokens: poolStateWithUnderlyingBalances.tokens.map(
+                (t) => t.underlyingToken ?? t,
+            ),
+        };
+
+        ({ bptAmount } = calculateProportionalAmounts(
+            poolStateWithBalances,
+            input.referenceAmount,
+        ));
+    }
+    return bptAmount;
+};