From a2f14a612208cea6a0e40776bbc2dd25908dedca Mon Sep 17 00:00:00 2001
From: johngrantuk <johngrantuk@googlemail.com>
Date: Thu, 25 Jul 2024 15:05:01 +0100
Subject: [PATCH] feat: Add dynamicFee support.

---
 typescript/src/hooks/types.ts            |   2 +-
 typescript/src/vault/vault.ts            |  13 ++-
 typescript/test/hooks/dynamicFee.test.ts | 111 +++++++++++++++++++++++
 3 files changed, 120 insertions(+), 6 deletions(-)
 create mode 100644 typescript/test/hooks/dynamicFee.test.ts

diff --git a/typescript/src/hooks/types.ts b/typescript/src/hooks/types.ts
index 9e24e6c..e3c5d62 100644
--- a/typescript/src/hooks/types.ts
+++ b/typescript/src/hooks/types.ts
@@ -66,8 +66,8 @@ export interface HookBase {
     };
     onComputeDynamicSwapFee(
         params: SwapInput,
-        poolAddress: string,
         staticSwapFeePercentage: bigint,
+        hookState: HookState | unknown,
     ): { success: boolean; dynamicSwapFee: bigint };
 }
 
diff --git a/typescript/src/vault/vault.ts b/typescript/src/vault/vault.ts
index f03bada..ab63aad 100644
--- a/typescript/src/vault/vault.ts
+++ b/typescript/src/vault/vault.ts
@@ -120,11 +120,14 @@ export class Vault {
             );
         }
 
-        // hook: dynamicSwapFee
+        let swapFee = poolState.swapFee;
         if (hook.shouldCallComputeDynamicSwapFee) {
-            throw new Error(
-                'Hook Unsupported: shouldCallComputeDynamicSwapFee',
+            const { success, dynamicSwapFee } = hook.onComputeDynamicSwapFee(
+                input,
+                poolState.swapFee,
+                hookState,
             );
+            if (success) swapFee = dynamicSwapFee;
         }
 
         // _swap()
@@ -140,13 +143,13 @@ export class Vault {
 
         // Set swapFeeAmountScaled18 based on the amountCalculated.
         let swapFeeAmountScaled18 = 0n;
-        if (poolState.swapFee > 0) {
+        if (swapFee > 0) {
             // Swap fee is always a percentage of the amountCalculated. On ExactIn, subtract it from the calculated
             // amountOut. On ExactOut, add it to the calculated amountIn.
             // Round up to avoid losses during precision loss.
             swapFeeAmountScaled18 = MathSol.mulUpFixed(
                 amountCalculatedScaled18,
-                poolState.swapFee,
+                swapFee,
             );
         }
 
diff --git a/typescript/test/hooks/dynamicFee.test.ts b/typescript/test/hooks/dynamicFee.test.ts
new file mode 100644
index 0000000..d9f5be3
--- /dev/null
+++ b/typescript/test/hooks/dynamicFee.test.ts
@@ -0,0 +1,111 @@
+// pnpm test -- dynamicFee.test.ts
+import { describe, expect, test } from 'vitest';
+import { SwapInput, SwapKind, Vault, Weighted } from '../../src';
+import { HookBase, HookState } from '@/hooks/types';
+
+const pool = {
+    poolType: 'CustomPool',
+    hookType: 'CustomHook',
+    chainId: '11155111',
+    blockNumber: '5955145',
+    poolAddress: '0xb2456a6f51530053bc41b0ee700fe6a2c37282e8',
+    tokens: [
+        '0x7b79995e5f793A07Bc00c21412e50Ecae098E7f9',
+        '0xb19382073c7A0aDdbb56Ac6AF1808Fa49e377B75',
+    ],
+    scalingFactors: [1000000000000000000n, 1000000000000000000n],
+    weights: [500000000000000000n, 500000000000000000n],
+    swapFee: 0n,
+
+    aggregateSwapFee: 500000000000000000n,
+    balancesLiveScaled18: [2000000000000000000n, 2000000000000000000n],
+    tokenRates: [1000000000000000000n, 1000000000000000000n],
+    totalSupply: 1000000000000000000n,
+};
+
+const swapInput = {
+    swapKind: SwapKind.GivenIn,
+    amountRaw: 100000000000000n,
+    tokenIn: pool.tokens[0],
+    tokenOut: pool.tokens[1],
+};
+
+describe('hook - dynamicFee', () => {
+    const vault = new Vault({
+        customPoolClasses: {
+            CustomPool: CustomPool,
+        },
+        customHookClasses: {
+            CustomHook: CustomHook,
+        },
+    });
+
+    test('should use dynamicFee set by hook', () => {
+        /*
+            hook state is used to pass new swapFee which gives expected swap result
+        */
+        const inputHookState = {
+            newSwapFee: 100000000000000000n,
+        };
+        const test = vault.swap(swapInput, pool, inputHookState);
+        expect(test).to.eq(89995500224987n);
+    });
+});
+
+class CustomPool extends Weighted {}
+
+class CustomHook implements HookBase {
+    public shouldCallComputeDynamicSwapFee = true;
+    public shouldCallBeforeSwap = false;
+    public shouldCallAfterSwap = false;
+    public shouldCallBeforeAddLiquidity = false;
+    public shouldCallAfterAddLiquidity = false;
+    public shouldCallBeforeRemoveLiquidity = false;
+    public shouldCallAfterRemoveLiquidity = false;
+    public enableHookAdjustedAmounts = false;
+
+    onBeforeAddLiquidity() {
+        return { success: false, hookAdjustedBalancesScaled18: [] };
+    }
+    onAfterAddLiquidity() {
+        return { success: false, hookAdjustedAmountsInRaw: [] };
+    }
+    onBeforeRemoveLiquidity() {
+        return { success: false, hookAdjustedBalancesScaled18: [] };
+    }
+    onAfterRemoveLiquidity() {
+        return {
+            success: false,
+            hookAdjustedAmountsOutRaw: [],
+        };
+    }
+    onBeforeSwap() {
+        return { success: false, hookAdjustedBalancesScaled18: [] };
+    }
+    onAfterSwap() {
+        return { success: false, hookAdjustedAmountCalculatedRaw: 0n };
+    }
+    onComputeDynamicSwapFee(
+        params: SwapInput,
+        staticSwapFeePercentage: bigint,
+        hookState: HookState | unknown,
+    ) {
+        if (
+            !(
+                typeof hookState === 'object' &&
+                hookState !== null &&
+                'newSwapFee' in hookState
+            )
+        )
+            throw new Error('Unexpected hookState');
+        expect(params.swapKind).to.eq(swapInput.swapKind);
+        expect(params.tokenIn).to.eq(swapInput.tokenIn);
+        expect(params.tokenOut).to.eq(swapInput.tokenOut);
+        expect(params.amountRaw).to.eq(swapInput.amountRaw);
+        expect(staticSwapFeePercentage).to.eq(pool.swapFee);
+        return {
+            success: true,
+            dynamicSwapFee: hookState.newSwapFee as bigint,
+        };
+    }
+}