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/getBoostedPoolStateWithBalancesV3.ts

@@ -0,0 +1,139 @@
+import {
+    createPublicClient,
+    erc4626Abi,
+    formatEther,
+    formatUnits,
+    http,
+    PublicClient,
+} from 'viem';
+
+import { HumanAmount } from '@/data';
+import { CHAINS, VAULT_V3 } from '@/utils';
+
+import {
+    PoolStateWithUnderlyingBalances,
+    PoolStateWithUnderlyings,
+} from '../types';
+import { vaultExtensionAbi_V3 } from '@/abi';
+import { getSortedTokens } from './getSortedTokens';
+import { TokenAmount } from '../tokenAmount';
+
+export const getBoostedPoolStateWithBalancesV3 = async (
+    poolState: PoolStateWithUnderlyings,
+    chainId: number,
+    rpcUrl: string,
+): Promise<PoolStateWithUnderlyingBalances> => {
+    const publicClient = createPublicClient({
+        transport: http(rpcUrl),
+        chain: CHAINS[chainId],
+    });
+
+    const { tokenBalances, totalShares } = await getTokenBalancesAndTotalShares(
+        chainId,
+        poolState,
+        publicClient,
+    );
+
+    const sortedTokens = [...poolState.tokens].sort(
+        (a, b) => a.index - b.index,
+    );
+
+    const underlyingBalances = await getUnderlyingBalances(
+        sortedTokens,
+        tokenBalances,
+        publicClient,
+    );
+
+    const poolStateWithUnderlyingBalances: PoolStateWithUnderlyingBalances = {
+        ...poolState,
+        tokens: sortedTokens.map((token, i) => ({
+            ...token,
+            balance: formatUnits(
+                tokenBalances[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: import('/Users/bruninho/Development/Balancer/b-sdk/src/entities/types').PoolTokenWithUnderlying[],
+    tokenBalances: TokenAmount[],
+    publicClient: PublicClient,
+) => {
+    const getUnderlyingBalancesContracts = sortedTokens
+        .filter((token) => token.underlyingToken !== null)
+        .map((token, i) => ({
+            address: token.address,
+            abi: erc4626Abi,
+            functionName: 'previewRedeem' as const,
+            args: [tokenBalances[i].amount] as const,
+        }));
+
+    const underlyingBalanceOutputs = await publicClient.multicall({
+        contracts: [...getUnderlyingBalancesContracts],
+    });
+    if (
+        underlyingBalanceOutputs.some((output) => output.status === 'failure')
+    ) {
+        throw new Error(
+            'Error: Unable to get underlying balances for v3 pool.',
+        );
+    }
+
+    const underlyingBalances = underlyingBalanceOutputs.map(
+        (output) => output.result as bigint,
+    );
+    return underlyingBalances;
+};
+
+const getTokenBalancesAndTotalShares = async (
+    chainId: number,
+    poolState: PoolStateWithUnderlyings,
+    publicClient: PublicClient,
+) => {
+    const totalSupplyContract = {
+        address: VAULT_V3[chainId],
+        abi: vaultExtensionAbi_V3,
+        functionName: 'totalSupply' as const,
+        args: [poolState.address] as const,
+    };
+    const getBalanceContracts = {
+        address: VAULT_V3[chainId],
+        abi: vaultExtensionAbi_V3,
+        functionName: 'getCurrentLiveBalances' as const,
+        args: [poolState.address] as const,
+    };
+
+    const outputs = await publicClient.multicall({
+        contracts: [totalSupplyContract, getBalanceContracts],
+    });
+
+    if (outputs.some((output) => output.status === 'failure')) {
+        throw new Error(
+            'Error: Unable to get pool state with balances for v3 pool.',
+        );
+    }
+
+    const totalShares = outputs[0].result as bigint;
+    const balancesScale18 = outputs[1].result as bigint[];
+
+    const poolTokens = getSortedTokens(poolState.tokens, chainId);
+    const tokenBalances = poolTokens.map((token, i) =>
+        TokenAmount.fromScale18Amount(token, balancesScale18[i]),
+    );
+
+    return { tokenBalances, totalShares };
+};

src/entities/utils/proportionalAmountsHelpers.ts

@@ -3,9 +3,11 @@ import { InputAmount } from '@/types';
 import { HumanAmount } from '@/data';
 import { 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 } from './getBoostedPoolStateWithBalancesV3';
+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 +129,41 @@ 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 (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;
+};

test/v3/addLiquidity/addLiquidity.integration.test.ts

@@ -438,6 +438,7 @@ describe('add liquidity test', () => {
                 );
             });
         });
+
         describe('add liquidity proportional', () => {
             let addLiquidityInput: AddLiquidityProportionalInput;
             beforeAll(() => {

test/v3/addLiquidityBoosted/addLiquidityBoosted.integration.test.ts

@@ -177,7 +177,7 @@ describe('Boosted AddLiquidity', () => {
             });
         });
         describe('add liquidity proportional', () => {
-            test('with tokens', async () => {
+            test('with bpt', async () => {
                 const addLiquidityProportionalInput: AddLiquidityBoostedInput =
                     {
                         chainId,
@@ -200,7 +200,7 @@ describe('Boosted AddLiquidity', () => {
                     };
 
                 const addLiquidityBuildCallOutput =
-                    await addLiquidityBoosted.buildCall(addLiquidityBuildInput);
+                    addLiquidityBoosted.buildCall(addLiquidityBuildInput);
 
                 const { transactionReceipt, balanceDeltas } =
                     await sendTransactionGetBalances(
@@ -261,6 +261,90 @@ describe('Boosted AddLiquidity', () => {
                     ),
                 );
             });
+            test('with reference token (non bpt)', async () => {
+                const addLiquidityProportionalInput: AddLiquidityBoostedInput =
+                    {
+                        chainId,
+                        rpcUrl,
+                        referenceAmount: {
+                            rawAmount: 481201n,
+                            decimals: 6,
+                            address: USDC.address,
+                        },
+                        kind: AddLiquidityKind.Proportional,
+                    };
+                const addLiquidityQueryOutput = await addLiquidityBoosted.query(
+                    addLiquidityProportionalInput,
+                    boostedPool_USDC_USDT,
+                );
+                const addLiquidityBuildInput: AddLiquidityBoostedBuildCallInput =
+                    {
+                        ...addLiquidityQueryOutput,
+                        slippage: Slippage.fromPercentage('1'),
+                    };
+
+                const addLiquidityBuildCallOutput =
+                    addLiquidityBoosted.buildCall(addLiquidityBuildInput);
+
+                const { transactionReceipt, balanceDeltas } =
+                    await sendTransactionGetBalances(
+                        [
+                            addLiquidityQueryOutput.bptOut.token.address,
+                            USDC.address as `0x${string}`,
+                            USDT.address as `0x${string}`,
+                        ],
+                        client,
+                        testAddress,
+                        addLiquidityBuildCallOutput.to, //
+                        addLiquidityBuildCallOutput.callData,
+                    );
+
+                expect(transactionReceipt.status).to.eq('success');
+
+                addLiquidityQueryOutput.amountsIn.map((a) => {
+                    expect(a.amount > 0n).to.be.true;
+                });
+
+                const expectedDeltas = [
+                    addLiquidityQueryOutput.bptOut.amount,
+                    ...addLiquidityQueryOutput.amountsIn.map(
+                        (tokenAmount) => tokenAmount.amount,
+                    ),
+                ];
+                expect(balanceDeltas).to.deep.eq(expectedDeltas);
+
+                const slippageAdjustedQueryInput =
+                    addLiquidityQueryOutput.amountsIn.map((amountsIn) => {
+                        return Slippage.fromPercentage('1').applyTo(
+                            amountsIn.amount,
+                            1,
+                        );
+                    });
+                expect(
+                    addLiquidityBuildCallOutput.maxAmountsIn.map(
+                        (a) => a.amount,
+                    ),
+                ).to.deep.eq(slippageAdjustedQueryInput);
+
+                // make sure to pass Tokens in correct order. Same as poolTokens but as underlyings instead
+                assertTokenMatch(
+                    [
+                        new Token(
+                            111555111,
+                            USDC.address as Address,
+                            USDC.decimals,
+                        ),
+                        new Token(
+                            111555111,
+                            USDT.address as Address,
+                            USDT.decimals,
+                        ),
+                    ],
+                    addLiquidityBuildCallOutput.maxAmountsIn.map(
+                        (a) => a.token,
+                    ),
+                );
+            });
             test('with native', async () => {
                 // TODO
             });