Add proportional nested exit

src/entities/index.ts

@@ -2,6 +2,7 @@ export * from './encoders';
 export * from './join';
 export * from './exit';
 export * from './nestedJoin';
+export * from './nestedExit';
 export * from './path';
 export * from './swap';
 export * from './slippage';

src/entities/nestedExit/doQueryNestedExit.ts

@@ -0,0 +1,42 @@
+import {
+    createPublicClient,
+    decodeAbiParameters,
+    decodeFunctionResult,
+    http,
+} from 'viem';
+import { Address, Hex } from '../../types';
+import { BALANCER_RELAYER, CHAINS } from '../../utils';
+import { balancerRelayerAbi } from '../../abi';
+
+export const doQueryNestedExit = async (
+    chainId: number,
+    rpcUrl: string,
+    accountAddress: Address,
+    encodedMulticall: Hex,
+    tokensOutLength: number,
+): Promise<bigint[]> => {
+    const client = createPublicClient({
+        transport: http(rpcUrl),
+        chain: CHAINS[chainId],
+    });
+
+    const { data } = await client.call({
+        account: accountAddress,
+        to: BALANCER_RELAYER[chainId],
+        data: encodedMulticall,
+    });
+
+    const result = decodeFunctionResult({
+        abi: balancerRelayerAbi,
+        functionName: 'vaultActionsQueryMulticall',
+        data: data as Hex,
+    });
+
+    const resultsToPeek = result.slice(result.length - tokensOutLength);
+
+    const peekedValues = resultsToPeek.map(
+        (r) => decodeAbiParameters([{ type: 'uint256' }], r)[0],
+    );
+
+    return peekedValues;
+};

src/entities/nestedExit/getNestedExitCalls.ts

@@ -0,0 +1,81 @@
+import { Token } from '../token';
+import { NestedExitInput, NestedExitCall } from './types';
+import { NestedPoolState } from '../types';
+import { TokenAmount } from '../tokenAmount';
+
+export const getNestedExitCalls = (
+    {
+        bptAmountIn,
+        chainId,
+        accountAddress,
+        useNativeAssetAsWrappedAmountOut = false,
+        toInternalBalance = false,
+    }: NestedExitInput,
+    { pools }: NestedPoolState,
+): { bptAmountIn: TokenAmount; calls: NestedExitCall[] } => {
+    /**
+     * Overall logic to build sequence of join calls:
+     * 1. Go from top pool to bottom filling out input amounts and output refs
+     * 2. Inputs will be bptAmountIn provided or output of the previous level
+     * 3. Output at bottom level is the amountsOut
+     */
+
+    // sort pools by descending level
+    const poolsSortedByLevel = pools.sort((a, b) => b.level - a.level);
+
+    const calls: NestedExitCall[] = [];
+    for (const pool of poolsSortedByLevel) {
+        const sortedTokens = pool.tokens
+            .sort((a, b) => a.index - b.index)
+            .map((t) => new Token(chainId, t.address, t.decimals));
+
+        // const sortedTokensWithoutBpt = sortedTokens.filter(
+        //     (t) => !t.isSameAddress(pool.address),
+        // );
+
+        const upperLevelCall = calls.find((call) =>
+            call.sortedTokens
+                .map((token) => token.address)
+                .includes(pool.address),
+        );
+        calls.push({
+            chainId: chainId,
+            useNativeAssetAsWrappedAmountOut,
+            sortedTokens,
+            poolId: pool.id,
+            poolType: pool.type,
+            kind: pool.type === 'ComposableStable' ? 3 : 0, // enum PoolKind { WEIGHTED, LEGACY_STABLE, COMPOSABLE_STABLE, COMPOSABLE_STABLE_V2 }
+            sender: accountAddress,
+            recipient: accountAddress,
+            bptAmountIn:
+                upperLevelCall === undefined
+                    ? {
+                          amount: bptAmountIn,
+                          isRef: false,
+                      }
+                    : {
+                          amount: upperLevelCall.outputReferenceKeys[
+                              upperLevelCall.sortedTokens
+                                  .map((token) => token.address)
+                                  .indexOf(pool.address)
+                          ],
+                          isRef: true,
+                      },
+            minAmountsOut: Array(sortedTokens.length).fill(0n),
+            toInternalBalance,
+            // TODO: previous implementation of nested exit didn't add an outputReferenceKey for the BPT token,
+            // but if I remove it from here, peek logic fails. Need to investigate why.
+            // Once we figure this out, we should be able to replace sortedTokens by sortedTokensWithoutBpt
+            outputReferenceKeys: sortedTokens.map(
+                (token) =>
+                    100n +
+                    BigInt(poolsSortedByLevel.indexOf(pool)) * 10n +
+                    BigInt(sortedTokens.indexOf(token)),
+            ),
+        });
+    }
+
+    const bptIn = new Token(chainId, poolsSortedByLevel[0].address, 18);
+    const _bptAmountIn = TokenAmount.fromRawAmount(bptIn, bptAmountIn);
+    return { calls, bptAmountIn: _bptAmountIn };
+};

src/entities/nestedExit/index.ts

@@ -0,0 +1,159 @@
+import { encodeFunctionData } from 'viem';
+import { Address, Hex } from '../../types';
+import { Token } from '../token';
+import { BALANCER_RELAYER, getPoolAddress } from '../../utils';
+import { Relayer } from '../relayer';
+import { TokenAmount } from '../tokenAmount';
+import { balancerRelayerAbi, bathcRelayerLibraryAbi } from '../../abi';
+import {
+    NestedExitInput,
+    NestedExitQueryResult,
+    NestedExitCallInput,
+} from './types';
+import { NestedPoolState } from '../types';
+import { doQueryNestedExit } from './doQueryNestedExit';
+import { getNestedExitCalls } from './getNestedExitCalls';
+import { parseNestedExitCall } from './parseNestedExitCall';
+
+export class NestedExit {
+    async query(
+        input: NestedExitInput,
+        nestedPoolState: NestedPoolState,
+    ): Promise<NestedExitQueryResult> {
+        const { calls, bptAmountIn } = getNestedExitCalls(
+            input,
+            nestedPoolState,
+        );
+
+        const parsedCalls = calls.map((call) => parseNestedExitCall(call));
+
+        const encodedCalls = parsedCalls.map((parsedCall) =>
+            encodeFunctionData({
+                abi: bathcRelayerLibraryAbi,
+                functionName: 'exitPool',
+                args: parsedCall.args,
+            }),
+        );
+
+        const tokensOut: Token[] = [];
+        const peekCalls: Hex[] = [];
+        calls.forEach((call) => {
+            call.outputReferenceKeys.forEach((opRefKey) => {
+                const tokenOut = call.sortedTokens[Number(opRefKey % 10n)];
+                const isTokenBeingUsedAsInput = calls.some(
+                    (_call) =>
+                        _call.bptAmountIn.isRef === true &&
+                        tokenOut.isSameAddress(getPoolAddress(_call.poolId)),
+                );
+
+                if (!isTokenBeingUsedAsInput) {
+                    tokensOut.push(tokenOut);
+                    peekCalls.push(
+                        Relayer.encodePeekChainedReferenceValue(
+                            Relayer.toChainedReference(opRefKey, false),
+                        ),
+                    );
+                }
+            });
+        });
+
+        // append peek calls to get amountsOut
+        encodedCalls.push(...peekCalls);
+
+        const encodedMulticall = encodeFunctionData({
+            abi: balancerRelayerAbi,
+            functionName: 'vaultActionsQueryMulticall',
+            args: [encodedCalls],
+        });
+
+        const peekedValues = await doQueryNestedExit(
+            input.chainId,
+            input.rpcUrl,
+            input.accountAddress,
+            encodedMulticall,
+            tokensOut.length,
+        );
+
+        console.log('peekedValues ', peekedValues);
+
+        const amountsOut = tokensOut.map((tokenOut, i) =>
+            TokenAmount.fromRawAmount(tokenOut, peekedValues[i]),
+        );
+
+        return { calls, bptAmountIn, amountsOut };
+    }
+
+    buildCall(input: NestedExitCallInput): {
+        call: Hex;
+        to: Address;
+        minAmountsOut: TokenAmount[];
+    } {
+        // apply slippage to amountsOut
+        const minAmountsOut = input.amountsOut.map((amountOut) =>
+            TokenAmount.fromRawAmount(
+                amountOut.token,
+                input.slippage.removeFrom(amountOut.amount),
+            ),
+        );
+
+        input.calls.forEach((call) => {
+            // update relevant calls with minAmountOut limits in place
+            minAmountsOut.forEach((minAmountOut, j) => {
+                const minAmountOutIndex = call.sortedTokens.findIndex((t) =>
+                    t.isSameAddress(minAmountOut.token.address),
+                );
+                if (minAmountOutIndex !== -1) {
+                    call.minAmountsOut[minAmountOutIndex] =
+                        minAmountsOut[j].amount;
+                }
+            });
+
+            // remove output reference key related to bpt token if present
+            // TODO: check why query/peek logic needs it
+            call.outputReferenceKeys = call.outputReferenceKeys.filter(
+                (_, i) => {
+                    const token = call.sortedTokens[i];
+                    const isBptToken = token.isSameAddress(
+                        getPoolAddress(call.poolId),
+                    );
+                    return !isBptToken;
+                },
+            );
+        });
+
+        const parsedCalls = input.calls.map((call) =>
+            parseNestedExitCall(call),
+        );
+
+        const encodedCalls = parsedCalls.map((parsedCall) =>
+            encodeFunctionData({
+                abi: bathcRelayerLibraryAbi,
+                functionName: 'exitPool',
+                args: parsedCall.args,
+            }),
+        );
+
+        // prepend relayer approval if provided
+        if (input.relayerApprovalSignature !== undefined) {
+            encodedCalls.unshift(
+                Relayer.encodeSetRelayerApproval(
+                    BALANCER_RELAYER[input.chainId],
+                    true,
+                    input.relayerApprovalSignature,
+                ),
+            );
+        }
+
+        const call = encodeFunctionData({
+            abi: balancerRelayerAbi,
+            functionName: 'multicall',
+            args: [encodedCalls],
+        });
+
+        return {
+            call,
+            to: BALANCER_RELAYER[input.chainId],
+            minAmountsOut,
+        };
+    }
+}

src/entities/nestedExit/parseNestedExitCall.ts

@@ -0,0 +1,69 @@
+import { Hex } from '../../types';
+import { WeightedEncoder } from '../encoders';
+import { ComposableStableEncoder } from '../encoders/composableStable';
+import { NestedExitCall } from './types';
+import { Relayer } from '../relayer';
+import { replaceWrapped } from '../utils/replaceWrapped';
+
+export const parseNestedExitCall = ({
+    useNativeAssetAsWrappedAmountOut,
+    chainId,
+    sortedTokens,
+    poolId,
+    poolType,
+    kind,
+    sender,
+    recipient,
+    bptAmountIn,
+    minAmountsOut,
+    toInternalBalance,
+    outputReferenceKeys,
+}: NestedExitCall) => {
+    // replace wrapped token with native asset if needed
+    let tokensOut = [...sortedTokens];
+    if (chainId && useNativeAssetAsWrappedAmountOut) {
+        tokensOut = replaceWrapped([...sortedTokens], chainId);
+    }
+
+    const _bptAmountIn = bptAmountIn.isRef
+        ? Relayer.toChainedReference(bptAmountIn.amount)
+        : bptAmountIn.amount;
+
+    let userData: Hex;
+    switch (poolType) {
+        case 'Weighted':
+            userData = WeightedEncoder.exitProportional(_bptAmountIn);
+            break;
+        case 'ComposableStable':
+            userData = ComposableStableEncoder.exitProportional(_bptAmountIn);
+            break;
+        default:
+            throw new Error('Unsupported pool type');
+    }
+
+    const outputReferences = outputReferenceKeys.map((k) => {
+        const tokenIndex = k % 10n;
+        return {
+            index: tokenIndex,
+            key: Relayer.toChainedReference(k),
+        };
+    });
+
+    const exitPoolRequest = {
+        assets: tokensOut.map((t) => t.address), // with BPT
+        minAmountsOut, // with BPT
+        userData, // wihtout BPT
+        toInternalBalance,
+    };
+
+    return {
+        args: [
+            poolId,
+            kind,
+            sender,
+            recipient,
+            exitPoolRequest,
+            outputReferences,
+        ] as const,
+    };
+};