infer unwrapWrapped for remove liquidity boosted

src/entities/addLiquidityBoosted/index.ts

@@ -55,9 +55,7 @@ export class AddLiquidityBoostedV3 {
         });
 
         const bptToken = new Token(input.chainId, poolState.address, 18);
-        const wrapUnderlying: boolean[] = new Array(
-            poolState.tokens.length,
-        ).fill(false);
+        const wrapUnderlying: boolean[] = [];
 
         let bptOut: TokenAmount;
         let amountsIn: TokenAmount[];
@@ -89,8 +87,8 @@ export class AddLiquidityBoostedV3 {
         switch (input.kind) {
             case AddLiquidityKind.Unbalanced: {
                 // Infer wrapUnderlying from token addreses provided by amountsIn
-                const tokensIn: ExtendedMinimalToken[] = input.amountsIn.map(
-                    (amountIn) => {
+                const tokensIn: ExtendedMinimalToken[] = input.amountsIn
+                    .map((amountIn) => {
                         const amountInAddress = amountIn.address.toLowerCase();
                         const token = poolStateTokenMap[amountInAddress];
                         if (!token) {
@@ -99,10 +97,11 @@ export class AddLiquidityBoostedV3 {
                             );
                         }
                         return token;
-                    },
-                );
+                    })
+                    .sort((a, b) => a.index - b.index); // sort by index so wrapUnderlying is in correct order
+
                 tokensIn.forEach((t) => {
-                    wrapUnderlying[t.index] = t.isUnderlyingToken;
+                    wrapUnderlying.push(t.isUnderlyingToken);
                 });
 
                 // It is allowed not not provide the same amount of TokenAmounts as inputs
@@ -131,14 +130,19 @@ export class AddLiquidityBoostedV3 {
             case AddLiquidityKind.Proportional: {
                 // User provides tokensIn addresses so we can infer if they need to be wrapped
 
-                input.tokensIn.forEach((t) => {
-                    const tokenWithDetails =
-                        poolStateTokenMap[t.toLowerCase() as Address];
-                    if (!tokenWithDetails) {
-                        throw new Error(`Invalid token address: ${t}`);
-                    }
-                    wrapUnderlying[tokenWithDetails.index] =
-                        tokenWithDetails.isUnderlyingToken;
+                const sortedTokensInWithDetails = input.tokensIn
+                    .map((t) => {
+                        const tokenWithDetails =
+                            poolStateTokenMap[t.toLowerCase() as Address];
+                        if (!tokenWithDetails) {
+                            throw new Error(`Invalid token address: ${t}`);
+                        }
+                        return tokenWithDetails;
+                    })
+                    .sort((a, b) => a.index - b.index);
+
+                sortedTokensInWithDetails.forEach((t) => {
+                    wrapUnderlying.push(t.isUnderlyingToken);
                 });
 
                 const bptAmount = await getBptAmountFromReferenceAmountBoosted(

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 { MinimalToken } from '@/data';
 
 export class RemoveLiquidityBoostedV3 implements RemoveLiquidityBase {
     private readonly inputValidator: InputValidator = new InputValidator();
@@ -39,6 +40,45 @@ export class RemoveLiquidityBoostedV3 implements RemoveLiquidityBase {
             type: 'Boosted',
         });
 
+        type ExtendedMinimalToken = MinimalToken & {
+            isUnderlyingToken: boolean;
+        };
+
+        // Build a useful map of all pool tokens with an isUnderlyingToken flag
+        const poolStateTokenMap: Record<Address, ExtendedMinimalToken> = {};
+        poolState.tokens.forEach((t) => {
+            const underlyingToken = t.underlyingToken;
+            poolStateTokenMap[t.address.toLowerCase()] = {
+                index: t.index,
+                decimals: t.decimals,
+                address: t.address,
+                isUnderlyingToken: false,
+            };
+            if (underlyingToken) {
+                poolStateTokenMap[underlyingToken.address.toLowerCase()] = {
+                    index: underlyingToken.index,
+                    decimals: underlyingToken.decimals,
+                    address: underlyingToken.address,
+                    isUnderlyingToken: true,
+                };
+            }
+        });
+
+        const sortedTokensOutWithDetails = input.tokensOut
+            .map((t) => {
+                const tokenWithDetails =
+                    poolStateTokenMap[t.toLowerCase() as Address];
+                if (!tokenWithDetails) {
+                    throw new Error(`Invalid token address: ${t}`);
+                }
+                return tokenWithDetails;
+            })
+            .sort((a, b) => a.index - b.index);
+
+        const unwrapWrapped: boolean[] = sortedTokensOutWithDetails.map(
+            (t) => t.isUnderlyingToken,
+        );
+
         const [tokensOut, underlyingAmountsOut] =
             await doRemoveLiquidityProportionalQuery(
                 input.rpcUrl,
@@ -47,7 +87,7 @@ export class RemoveLiquidityBoostedV3 implements RemoveLiquidityBase {
                 input.sender ?? zeroAddress,
                 input.userData ?? '0x',
                 poolState.address,
-                input.unwrapWrapped,
+                unwrapWrapped,
                 block,
             );
 
@@ -77,7 +117,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;

test/lib/utils/addLiquidityBoostedHelper.ts

@@ -38,7 +38,6 @@ export async function GetBoostedBpt(
         decimals: number;
         slot: number;
     }[],
-    wrapUnderlying: boolean[],
 ): Promise<bigint> {
     await setTokenBalances(
         client,
@@ -50,7 +49,6 @@ export async function GetBoostedBpt(
 
     const addLiquidityInput: AddLiquidityBoostedUnbalancedInput = {
         amountsIn,
-        wrapUnderlying,
         chainId,
         rpcUrl,
         kind: AddLiquidityKind.Unbalanced,

test/v3/removeLiquidityBoosted/removeLiquidityBoosted.integration.test.ts

@@ -140,7 +140,7 @@ describe('remove liquidity boosted proportional', () => {
                 decimals: 18,
                 address: boostedPool_USDC_USDT.address,
             },
-            unwrapWrapped: [true, false],
+            tokensOut: [USDC.address, stataUSDT.address],
             kind: RemoveLiquidityKind.Proportional,
         };
 
@@ -179,7 +179,7 @@ describe('remove liquidity boosted proportional', () => {
                         decimals: 18,
                         address: boostedPool_USDC_USDT.address,
                     },
-                    unwrapWrapped: [true, true],
+                    tokensOut: [USDC.address, USDT.address],
                     kind: RemoveLiquidityKind.Proportional,
                 };
 
@@ -271,7 +271,7 @@ describe('remove liquidity boosted proportional', () => {
                         decimals: 18,
                         address: boostedPool_USDC_USDT.address,
                     },
-                    unwrapWrapped: [true, false],
+                    tokensOut: [USDC.address, stataUSDT.address],
                     kind: RemoveLiquidityKind.Proportional,
                 };
 
@@ -364,7 +364,7 @@ describe('remove liquidity boosted proportional', () => {
                         decimals: 18,
                         address: boostedPool_USDC_USDT.address,
                     },
-                    unwrapWrapped: [true, true],
+                    tokensOut: [USDC.address, USDT.address],
                     kind: RemoveLiquidityKind.Proportional,
                     sender: testAddress,
                     userData: '0x123',

test/v3/removeLiquidityBoosted/removeLiquidityPartialBoosted.integration.test.ts

@@ -51,7 +51,6 @@ const parentBptToken = new Token(
 // These are the underlying tokens
 const usdtToken = new Token(chainId, USDT.address, USDT.decimals);
 const wethToken = new Token(chainId, WETH.address, WETH.decimals);
-const unwrapWrapped = [false, true]; // unwrapWrapped must match order of on chain state for pool tokens
 
 describe('V3 remove liquidity partial boosted', () => {
     let rpcUrl: string;
@@ -99,7 +98,6 @@ describe('V3 remove liquidity partial boosted', () => {
                     slot: WETH.slot as number,
                 },
             ],
-            unwrapWrapped,
         );
 
         removeLiquidityInput = {
@@ -110,7 +108,7 @@ describe('V3 remove liquidity partial boosted', () => {
                 decimals: 18,
                 rawAmount: bptAmount,
             },
-            unwrapWrapped,
+            tokensOut: [USDT.address, WETH.address],
             kind: RemoveLiquidityKind.Proportional,
         };