feat: Add support for after add hook.

typescript/src/hooks/types.ts

@@ -36,13 +36,12 @@ export interface HookBase {
         balancesScaled18: bigint[],
     ): boolean;
     onAfterAddLiquidity(
-        router: string,
-        pool: string,
         kind: AddKind,
         amountsInScaled18: bigint[],
         amountsInRaw: bigint[],
         bptAmountOut: bigint,
         balancesScaled18: bigint[],
+        hookState: HookState | unknown,
     ): { success: boolean; hookAdjustedAmountsInRaw: bigint[] };
     onBeforeRemoveLiquidity(
         router: string,

typescript/src/vault/basePoolMath.ts

@@ -23,7 +23,7 @@ export function computeAddLiquidityUnbalanced(
     // Create a new array to hold the updated balances after the addition.
     const newBalances: bigint[] = new Array(numTokens);
     // Create a new array to hold the swap fee amount for each token.
-    const swapFeeAmounts: bigint[] = new Array(numTokens);
+    const swapFeeAmounts: bigint[] = new Array(numTokens).fill(0n);
 
     // Loop through each token, updating the balance with the added amount.
     for (let index = 0; index < currentBalances.length; index++) {
@@ -51,13 +51,13 @@ export function computeAddLiquidityUnbalanced(
                 newBalances[index] -
                 MathSol.mulUpFixed(invariantRatio, currentBalances[index]);
             // Calculate fee amount
-            const swapFeeAmount = MathSol.mulUpFixed(
+            swapFeeAmounts[index] = MathSol.mulUpFixed(
                 taxableAmount,
                 swapFeePercentage,
             );
             // Subtract the fee from the new balance.
             // We are essentially imposing swap fees on non-proportional incoming amounts.
-            newBalances[index] = newBalances[index] - swapFeeAmount;
+            newBalances[index] = newBalances[index] - swapFeeAmounts[index];
         }
     }
 

typescript/src/vault/vault.ts

@@ -168,7 +168,7 @@ export class Vault {
     public addLiquidity(
         input: AddLiquidityInput,
         poolState: PoolState,
-        hookState?: HookState,
+        hookState?: HookState | unknown,
     ): { amountsIn: bigint[]; bptAmountOut: bigint } {
         if (poolState.poolType === 'Buffer')
             throw Error('Buffer pools do not support addLiquidity');
@@ -194,6 +194,8 @@ export class Vault {
 
         let amountsInScaled18: bigint[];
         let bptAmountOut: bigint;
+        let swapFeeAmountsScaled18: bigint[];
+
         if (input.kind === AddKind.UNBALANCED) {
             amountsInScaled18 = maxAmountsInScaled18;
             const computed = computeAddLiquidityUnbalanced(
@@ -205,6 +207,7 @@ export class Vault {
                     pool.computeInvariant(balancesLiveScaled18),
             );
             bptAmountOut = computed.bptAmountOut;
+            swapFeeAmountsScaled18 = computed.swapFeeAmounts;
         } else if (input.kind === AddKind.SINGLE_TOKEN_EXACT_OUT) {
             const tokenIndex = this._getSingleInputIndex(maxAmountsInScaled18);
             amountsInScaled18 = maxAmountsInScaled18;
@@ -223,21 +226,53 @@ export class Vault {
                     ),
             );
             amountsInScaled18[tokenIndex] = computed.amountInWithFee;
+            swapFeeAmountsScaled18 = computed.swapFeeAmounts;
         } else throw new Error('Unsupported AddLiquidity Kind');
 
         const amountsInRaw: bigint[] = new Array(poolState.tokens.length);
+        const updatedBalancesLiveScaled18 = [...poolState.balancesLiveScaled18];
         for (let i = 0; i < poolState.tokens.length; i++) {
             // amountsInRaw are amounts actually entering the Pool, so we round up.
             amountsInRaw[i] = this._toRawUndoRateRoundUp(
                 amountsInScaled18[i],
                 poolState.scalingFactors[i],
                 poolState.tokenRates[i],
             );
+
+            // A Pool's token balance always decreases after an exit
+            // Computes protocol and pool creator fee which is eventually taken from pool balance
+            const aggregateSwapFeeAmountScaled18 =
+                this._computeAndChargeAggregateSwapFees(
+                    swapFeeAmountsScaled18[i],
+                    poolState.aggregateSwapFee,
+                );
+
+            updatedBalancesLiveScaled18[i] =
+                poolState.balancesLiveScaled18[i] +
+                amountsInScaled18[i] -
+                aggregateSwapFeeAmountScaled18;
         }
 
-        // hook: shouldCallAfterAddLiquidity
         if (hook.shouldCallAfterAddLiquidity) {
-            throw new Error('Hook Unsupported: shouldCallAfterAddLiquidity');
+            const { success, hookAdjustedAmountsInRaw } =
+                hook.onAfterAddLiquidity(
+                    input.kind,
+                    amountsInScaled18,
+                    amountsInRaw,
+                    bptAmountOut,
+                    updatedBalancesLiveScaled18,
+                    hookState,
+                );
+
+            if (
+                success === false ||
+                hookAdjustedAmountsInRaw.length != amountsInRaw.length
+            ) {
+                throw new Error(
+                    `AfterAddLiquidityHookFailed ${poolState.poolType} ${poolState.hookType}`,
+                );
+            }
+            hookAdjustedAmountsInRaw.forEach((a, i) => (amountsInRaw[i] = a));
         }
 
         return {

typescript/test/hooks/afterAddLiquidity.test.ts

@@ -0,0 +1,151 @@
+// pnpm test -- afterAddLiquidity.test.ts
+import { describe, expect, test } from 'vitest';
+import { AddKind, Vault, Weighted } from '../../src';
+import { HookBase, HookState } from '@/hooks/types';
+
+const addLiquidityInput = {
+    pool: '0xb2456a6f51530053bc41b0ee700fe6a2c37282e8',
+    maxAmountsIn: [200000000000000000n, 100000000000000000n],
+    minBptAmountOut: 0n,
+    kind: AddKind.UNBALANCED,
+};
+
+class CustomPool extends Weighted {}
+
+const pool = {
+    poolType: 'CustomPool',
+    hookType: 'CustomHook',
+    chainId: '11155111',
+    blockNumber: '5955145',
+    poolAddress: '0xb2456a6f51530053bc41b0ee700fe6a2c37282e8',
+    tokens: [
+        '0x7b79995e5f793A07Bc00c21412e50Ecae098E7f9',
+        '0xb19382073c7A0aDdbb56Ac6AF1808Fa49e377B75',
+    ],
+    scalingFactors: [1000000000000000000n, 1000000000000000000n],
+    weights: [500000000000000000n, 500000000000000000n],
+    swapFee: 100000000000000000n,
+    aggregateSwapFee: 500000000000000000n,
+    balancesLiveScaled18: [1000000000000000000n, 1000000000000000000n],
+    tokenRates: [1000000000000000000n, 1000000000000000000n],
+    totalSupply: 1000000000000000000n,
+};
+
+describe('hook - afterAddLiquidity', () => {
+    const vault = new Vault({
+        customPoolClasses: {
+            CustomPool: CustomPool,
+        },
+        customHookClasses: {
+            CustomHook: CustomHook,
+        },
+    });
+
+    test('aggregateSwapFee of 0 should not take any protocol fees from updated balances', () => {
+        // hook state is used to pass expected value to tests
+        const inputHookState = {
+            expectedBalancesLiveScaled18: [
+                1200000000000000000n,
+                1100000000000000000n,
+            ],
+        };
+        const test = vault.addLiquidity(
+            addLiquidityInput,
+            {
+                ...pool,
+                aggregateSwapFee: 0n,
+            },
+            inputHookState,
+        );
+        // Hook adds 1n to amountsIn
+        expect(test.amountsIn).to.deep.eq([
+            200000000000000001n,
+            100000000000000001n,
+        ]);
+        expect(test.bptAmountOut).to.deep.eq(146464294351915965n);
+    });
+
+    test('aggregateSwapFee of 50% should take half of remaining', () => {
+        /*
+            hook state is used to pass expected value to tests
+            aggregate fee amount is 2554373534619714n which is deducted from amount in
+        */
+        const inputHookState = {
+            expectedBalancesLiveScaled18: [
+                1197445626465380286n,
+                1100000000000000000n,
+            ],
+        };
+        const test = vault.addLiquidity(
+            addLiquidityInput,
+            pool,
+            inputHookState,
+        );
+        expect(test.amountsIn).to.deep.eq([
+            200000000000000001n,
+            100000000000000001n,
+        ]);
+        expect(test.bptAmountOut).to.deep.eq(146464294351915965n);
+    });
+});
+
+class CustomHook implements HookBase {
+    public shouldCallComputeDynamicSwapFee = false;
+    public shouldCallBeforeSwap = false;
+    public shouldCallAfterSwap = false;
+    public shouldCallBeforeAddLiquidity = false;
+    public shouldCallAfterAddLiquidity = true;
+    public shouldCallBeforeRemoveLiquidity = false;
+    public shouldCallAfterRemoveLiquidity = false;
+    public enableHookAdjustedAmounts = true;
+
+    onBeforeAddLiquidity() {
+        return false;
+    }
+    onAfterAddLiquidity(
+        kind: AddKind,
+        amountsInScaled18: bigint[],
+        amountsInRaw: bigint[],
+        bptAmountOut: bigint,
+        balancesScaled18: bigint[],
+        hookState: HookState | unknown,
+    ) {
+        if (
+            !(
+                typeof hookState === 'object' &&
+                hookState !== null &&
+                'expectedBalancesLiveScaled18' in hookState
+            )
+        )
+            throw new Error('Unexpected hookState');
+        expect(kind).to.eq(addLiquidityInput.kind);
+        expect(bptAmountOut).to.eq(146464294351915965n);
+        expect(amountsInScaled18).to.deep.eq(addLiquidityInput.maxAmountsIn);
+        expect(amountsInRaw).to.deep.eq(addLiquidityInput.maxAmountsIn);
+        expect(balancesScaled18).to.deep.eq(
+            hookState.expectedBalancesLiveScaled18,
+        );
+        return {
+            success: true,
+            hookAdjustedAmountsInRaw: [
+                amountsInRaw[0] + 1n,
+                amountsInRaw[1] + 1n,
+            ],
+        };
+    }
+    onBeforeRemoveLiquidity() {
+        return false;
+    }
+    onAfterRemoveLiquidity() {
+        return { success: false, hookAdjustedAmountsOutRaw: [] };
+    }
+    onBeforeSwap() {
+        return false;
+    }
+    onAfterSwap() {
+        return { success: false, hookAdjustedAmountCalculatedRaw: 0n };
+    }
+    onComputeDynamicSwapFee() {
+        return { success: false, dynamicSwapFee: 0n };
+    }
+}