New composite router for boosted pool operations

src/entities/addLiquidityBoosted/index.ts

@@ -2,27 +2,21 @@
 // available:
 // 1. Unbalanced - addLiquidityUnbalancedToERC4626Pool
 // 2. Proportional - addLiquidityProportionalToERC4626Pool
-import { encodeFunctionData, zeroAddress } from 'viem';
+import { encodeFunctionData, zeroAddress, Address } from 'viem';
 import { TokenAmount } from '@/entities/tokenAmount';
-
 import { Permit2 } from '@/entities/permit2Helper';
-
 import { getAmountsCall } from '../addLiquidity/helpers';
-
 import { PoolStateWithUnderlyings } from '@/entities/types';
-
 import {
     getAmounts,
     getBptAmountFromReferenceAmountBoosted,
     getSortedTokens,
     getValue,
 } from '@/entities/utils';
-
 import {
     AddLiquidityBuildCallOutput,
     AddLiquidityKind,
 } from '../addLiquidity/types';
-
 import { doAddLiquidityUnbalancedQuery } from './doAddLiquidityUnbalancedQuery';
 import { doAddLiquidityProportionalQuery } from './doAddLiquidityPropotionalQuery';
 import { Token } from '../token';
@@ -31,14 +25,17 @@ import {
     balancerCompositeLiquidityRouterBoostedAbi,
     balancerRouterAbi,
 } from '@/abi';
-
 import { Hex } from '@/types';
 import {
     AddLiquidityBoostedBuildCallInput,
     AddLiquidityBoostedInput,
     AddLiquidityBoostedQueryOutput,
 } from './types';
 import { InputValidator } from '../inputValidator/inputValidator';
+import {
+    buildPoolStateTokenMap,
+    MinimalTokenWithIsUnderlyingFlag,
+} from '@/entities/utils';
 
 export class AddLiquidityBoostedV3 {
     private readonly inputValidator: InputValidator = new InputValidator();
@@ -54,41 +51,52 @@ export class AddLiquidityBoostedV3 {
         });
 
         const bptToken = new Token(input.chainId, poolState.address, 18);
+        const wrapUnderlying: boolean[] = new Array(
+            poolState.tokens.length,
+        ).fill(false);
 
         let bptOut: TokenAmount;
         let amountsIn: TokenAmount[];
 
+        const poolStateTokenMap = buildPoolStateTokenMap(poolState);
+
         switch (input.kind) {
             case AddLiquidityKind.Unbalanced: {
-                // Use poolState to add token index to amountsIn
-                const tokensIn = input.amountsIn.map((amountIn) => {
-                    const tokenIn = poolState.tokens.find((t) => {
+                // Use amountsIn provided by use to infer if token should be wrapped
+                const tokensIn: MinimalTokenWithIsUnderlyingFlag[] =
+                    input.amountsIn.map((amountIn) => {
                         const amountInAddress = amountIn.address.toLowerCase();
-                        return (
-                            t.address.toLowerCase() === amountInAddress ||
-                            (t.underlyingToken &&
-                                t.underlyingToken.address.toLowerCase() ===
-                                    amountInAddress)
-                        );
+                        const token = poolStateTokenMap[amountInAddress];
+                        if (!token) {
+                            throw new Error(
+                                `Token not found in poolState: ${amountInAddress}`,
+                            );
+                        }
+                        return token;
                     });
 
-                    if (!tokenIn) {
-                        throw new Error(
-                            `Token not found in poolState: ${amountIn.address}`,
-                        );
-                    }
-
-                    const matchedToken =
-                        tokenIn.underlyingToken?.address.toLowerCase() ===
-                        amountIn.address.toLowerCase()
-                            ? tokenIn.underlyingToken
-                            : tokenIn;
+                // if user provides fewer than the number of pool tokens,
+                // fill remaining indexes because composite router requires
+                // length of pool tokens to match maxAmountsIn and wrapUnderlying
+                if (tokensIn.length < poolState.tokens.length) {
+                    const existingIndices = new Set(
+                        tokensIn.map((t) => t.index),
+                    );
+                    poolState.tokens.forEach((poolToken) => {
+                        if (!existingIndices.has(poolToken.index)) {
+                            tokensIn.push({
+                                index: poolToken.index,
+                                decimals: poolToken.decimals,
+                                address: poolToken.address,
+                                isUnderlyingToken: false,
+                            });
+                        }
+                    });
+                }
 
-                    return {
-                        address: matchedToken.address,
-                        decimals: matchedToken.decimals,
-                        index: matchedToken.index,
-                    };
+                // wrap if token is underlying
+                tokensIn.forEach((t) => {
+                    wrapUnderlying[t.index] = t.isUnderlyingToken;
                 });
 
                 // It is allowed not not provide the same amount of TokenAmounts as inputs
@@ -103,7 +111,7 @@ export class AddLiquidityBoostedV3 {
                     input.sender ?? zeroAddress,
                     input.userData ?? '0x',
                     poolState.address,
-                    input.wrapUnderlying,
+                    wrapUnderlying,
                     maxAmountsIn,
                     block,
                 );
@@ -112,12 +120,24 @@ export class AddLiquidityBoostedV3 {
                 amountsIn = sortedTokens.map((t, i) =>
                     TokenAmount.fromRawAmount(t, maxAmountsIn[i]),
                 );
+
                 break;
             }
             case AddLiquidityKind.Proportional: {
+                // User provides addresses via input.tokensIn so we can infer if they need to be wrapped
+                input.tokensIn.forEach((t) => {
+                    const tokenIn =
+                        poolStateTokenMap[t.toLowerCase() as Address];
+                    if (!tokenIn) {
+                        throw new Error(`Invalid token address: ${t}`);
+                    }
+                    wrapUnderlying[tokenIn.index] = tokenIn.isUnderlyingToken;
+                });
+
                 const bptAmount = await getBptAmountFromReferenceAmountBoosted(
                     input,
                     poolState,
+                    wrapUnderlying,
                 );
 
                 const [tokensIn, exactAmountsInNumbers] =
@@ -128,26 +148,13 @@ export class AddLiquidityBoostedV3 {
                         input.userData ?? '0x',
                         poolState.address,
                         bptAmount.rawAmount,
-                        input.wrapUnderlying,
+                        wrapUnderlying,
                         block,
                     );
 
-                amountsIn = tokensIn.map((tokenInAddress, i) => {
-                    const decimals = poolState.tokens.find((t) => {
-                        const tokenAddress = tokenInAddress.toLowerCase();
-                        return (
-                            t.address.toLowerCase() === tokenAddress ||
-                            (t.underlyingToken &&
-                                t.underlyingToken.address.toLowerCase() ===
-                                    tokenAddress)
-                        );
-                    })?.decimals;
-
-                    if (!decimals)
-                        throw new Error(
-                            `Token decimals missing for ${tokenInAddress}`,
-                        );
-
+                amountsIn = tokensIn.map((t, i) => {
+                    const tokenInAddress = t.toLowerCase() as `0x${string}`;
+                    const { decimals } = poolStateTokenMap[tokenInAddress];
                     const token = new Token(
                         input.chainId,
                         tokenInAddress,
@@ -172,7 +179,7 @@ export class AddLiquidityBoostedV3 {
             poolId: poolState.id,
             poolType: poolState.type,
             addLiquidityKind: input.kind,
-            wrapUnderlying: input.wrapUnderlying,
+            wrapUnderlying,
             bptOut,
             amountsIn,
             chainId: input.chainId,

src/entities/addLiquidityBoosted/types.ts

@@ -7,9 +7,9 @@ import { Slippage } from '../slippage';
 export type AddLiquidityBoostedProportionalInput = {
     chainId: number;
     rpcUrl: string;
+    tokensIn: Address[];
     referenceAmount: InputAmount;
     kind: AddLiquidityKind.Proportional;
-    wrapUnderlying: boolean[];
     sender?: Address;
     userData?: Hex;
 };
@@ -19,7 +19,6 @@ export type AddLiquidityBoostedUnbalancedInput = {
     rpcUrl: string;
     amountsIn: InputAmount[];
     kind: AddLiquidityKind.Unbalanced;
-    wrapUnderlying: boolean[];
     sender?: Address;
     userData?: Hex;
 };

src/entities/addLiquidityNested/addLiquidityNestedV3/types.ts

@@ -5,7 +5,6 @@ import { Address, Hex } from 'viem';
 import { AddLiquidityNestedBaseInput } from '../types';
 
 export type AddLiquidityNestedInputV3 = AddLiquidityNestedBaseInput & {
-    wrapUnderlying?: boolean[];
     sender?: Address;
     userData?: Hex;
 };

src/entities/priceImpact/addLiquidityNested.ts

@@ -64,14 +64,13 @@ export async function addLiquidityNested(
             });
         }
         let addResult: AddResult;
-        if (isBoostedPool && 'wrapUnderlying' in input && input.wrapUnderlying)
+        if (isBoostedPool)
             addResult = await getAddBoostedUnbalancedResult(
                 addLiquidityBoosted,
                 input.chainId,
                 input.rpcUrl,
                 pool as NestedPoolV3,
                 amountsIn,
-                input.wrapUnderlying,
             );
         else
             addResult = await getAddUnbalancedResult(
@@ -129,14 +128,12 @@ async function getAddBoostedUnbalancedResult(
     rpcUrl: string,
     pool: NestedPoolV3,
     amountsIn: InputAmount[],
-    wrapUnderlying: boolean[],
 ): Promise<AddResult> {
     const addLiquidityInput: AddLiquidityBoostedUnbalancedInput = {
         chainId,
         rpcUrl,
         amountsIn,
         kind: AddLiquidityKind.Unbalanced,
-        wrapUnderlying,
     };
 
     const priceImpactAmount = await addLiquidityUnbalancedBoosted(

src/entities/priceImpact/addLiquidityUnbalancedBoosted.ts

@@ -66,7 +66,7 @@ export async function addLiquidityUnbalancedBoosted(
         chainId: input.chainId,
         rpcUrl: input.rpcUrl,
         bptIn: bptOut.toInputAmount(),
-        unwrapWrapped: input.wrapUnderlying,
+        tokensOut: poolTokens.map((t) => t.address),
         kind: RemoveLiquidityKind.Proportional,
     };
     const { amountsOut } = await removeLiquidity.query(

src/entities/removeLiquidityBoosted/index.ts

@@ -1,4 +1,4 @@
-import { encodeFunctionData, zeroAddress } from 'viem';
+import { Address, encodeFunctionData, zeroAddress } from 'viem';
 
 import {
     RemoveLiquidityBase,
@@ -25,6 +25,7 @@ import {
     RemoveLiquidityBoostedQueryOutput,
 } from './types';
 import { InputValidator } from '../inputValidator/inputValidator';
+import { buildPoolStateTokenMap } from '@/entities/utils';
 
 export class RemoveLiquidityBoostedV3 implements RemoveLiquidityBase {
     private readonly inputValidator: InputValidator = new InputValidator();
@@ -39,6 +40,20 @@ export class RemoveLiquidityBoostedV3 implements RemoveLiquidityBase {
             type: 'Boosted',
         });
 
+        const poolStateTokenMap = buildPoolStateTokenMap(poolState);
+
+        // Infer if token should be unwrapped using the tokensOut provided by user
+        const unwrapWrapped = input.tokensOut
+            .map((t) => {
+                const tokenOut = poolStateTokenMap[t.toLowerCase() as Address];
+                if (!tokenOut) {
+                    throw new Error(`Invalid token address: ${t}`);
+                }
+                return tokenOut;
+            })
+            .sort((a, b) => a.index - b.index) // sort by index to match the order of the pool tokens
+            .map((t) => t.isUnderlyingToken);
+
         const [tokensOut, underlyingAmountsOut] =
             await doRemoveLiquidityProportionalQuery(
                 input.rpcUrl,
@@ -47,25 +62,14 @@ export class RemoveLiquidityBoostedV3 implements RemoveLiquidityBase {
                 input.sender ?? zeroAddress,
                 input.userData ?? '0x',
                 poolState.address,
-                input.unwrapWrapped,
+                unwrapWrapped,
                 block,
             );
 
-        // tokens out can be underlying or yield-bearing variant
         const amountsOut = underlyingAmountsOut.map((amount, i) => {
             const tokenOut = tokensOut[i];
-
-            const decimals = poolState.tokens.find((t) => {
-                return (
-                    t.address.toLowerCase() === tokenOut.toLowerCase() ||
-                    (t.underlyingToken &&
-                        t.underlyingToken.address.toLowerCase() ===
-                            tokenOut.toLowerCase())
-                );
-            })?.decimals;
-
-            if (!decimals)
-                throw new Error(`Token decimals missing for ${tokenOut}`);
+            const { decimals } =
+                poolStateTokenMap[tokenOut.toLowerCase() as Address];
 
             const token = new Token(input.chainId, tokenOut, decimals);
             return TokenAmount.fromRawAmount(token, amount);
@@ -77,7 +81,7 @@ export class RemoveLiquidityBoostedV3 implements RemoveLiquidityBase {
             to: BALANCER_COMPOSITE_LIQUIDITY_ROUTER_BOOSTED[input.chainId],
             poolType: poolState.type,
             poolId: poolState.address,
-            unwrapWrapped: input.unwrapWrapped,
+            unwrapWrapped,
             removeLiquidityKind: RemoveLiquidityKind.Proportional,
             bptIn: TokenAmount.fromRawAmount(bptToken, input.bptIn.rawAmount),
             amountsOut,

src/entities/removeLiquidityBoosted/types.ts

@@ -8,7 +8,7 @@ export type RemoveLiquidityBoostedProportionalInput = {
     chainId: number;
     rpcUrl: string;
     bptIn: InputAmount;
-    unwrapWrapped: boolean[];
+    tokensOut: Address[];
     kind: RemoveLiquidityKind.Proportional;
     sender?: Address;
     userData?: Hex;