55 python hook tests

python/src/add_liquidity.py

@@ -12,7 +12,7 @@
 
 
 class Kind(Enum):
-    PROPORTIONAL = 0
+    UNBALANCED = 0
     SINGLE_TOKEN_EXACT_OUT = 1
 
 
@@ -29,13 +29,13 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_
     )
 
     updated_balances_live_scaled18 = pool_state["balancesLiveScaled18"][:]
-    if hook_class.shouldCallBeforeAddLiquidity:
+    if hook_class.should_call_before_add_liquidity:
         # Note - in SC balances and amounts are updated to reflect any rate change.
         # Daniel said we should not worry about this as any large rate changes
         # will mean something has gone wrong.
         # We do take into account and balance changes due
         # to hook using hookAdjustedBalancesScaled18.
-        hook_return = hook_class.onBeforeSwap(
+        hook_return = hook_class.on_before_add_liquidity(
             add_liquidity_input["kind"],
             add_liquidity_input["max_amounts_in_raw"],
             add_liquidity_input["min_bpt_amount_out_raw"],
@@ -44,10 +44,10 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_
         )
         if hook_return["success"] is False:
             raise SystemError("BeforeAddLiquidityHookFailed")
-        for i, a in enumerate(hook_return["hookAdjustedBalancesScaled18"]):
+        for i, a in enumerate(hook_return["hook_adjusted_balances_scaled18"]):
             updated_balances_live_scaled18[i] = a
 
-    if add_liquidity_input["kind"] == Kind.PROPORTIONAL.value:
+    if add_liquidity_input["kind"] == Kind.UNBALANCED.value:
         amounts_in_scaled18 = max_amounts_in_scaled18
         computed = compute_add_liquidity_unbalanced(
             updated_balances_live_scaled18,
@@ -98,24 +98,24 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_
         )
 
         # Update the balances with the incoming amounts and subtract the swap fees
-        updated_balances_live_scaled18[i] += (
-            amounts_in_scaled18[i] - aggregate_swap_fee_amount_scaled18
+        updated_balances_live_scaled18[i] = (
+            updated_balances_live_scaled18[i]
+            + amounts_in_scaled18[i]
+            - aggregate_swap_fee_amount_scaled18
         )
 
-    if hook_class.shouldCallAfterAddLiquidity:
+    if hook_class.should_call_after_add_liquidity:
         hook_return = hook_class.on_after_add_liquidity(
-            {
-                "kind": add_liquidity_input["kind"],
-                "amounts_in_scaled18": amounts_in_scaled18,
-                "amounts_in_raw": amounts_in_raw,
-                "bpt_amount_out": bpt_amount_out,
-                "balances_scaled_18": updated_balances_live_scaled18,
-                "hook_state": hook_state,
-            }
+            add_liquidity_input["kind"],
+            amounts_in_scaled18,
+            amounts_in_raw,
+            bpt_amount_out,
+            updated_balances_live_scaled18,
+            hook_state,
         )
 
         if hook_return["success"] is False or len(
-            hook_return["hookAdjustedAmountsInRaw"]
+            hook_return["hook_adjusted_amounts_in_raw"]
         ) is not len(amounts_in_raw):
             raise SystemError(
                 " AfterAddLiquidityHookFailed",
@@ -124,7 +124,7 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_
             )
 
         # If hook adjusted amounts is not enabled, ignore amounts returned by the hook
-        if hook_class.enableHookAdjustedAmounts:
+        if hook_class.enable_hook_adjusted_amounts:
             for i, a in enumerate(hook_return["hook_adjusted_amounts_in_raw"]):
                 amounts_in_raw[i] = a
 

python/src/hooks/default_hook.py

@@ -1,12 +1,12 @@
 class DefaultHook:
-    shouldCallComputeDynamicSwapFee = False
-    shouldCallBeforeSwap = False
-    shouldCallAfterSwap = False
-    shouldCallBeforeAddLiquidity = False
-    shouldCallAfterAddLiquidity = False
-    shouldCallBeforeRemoveLiquidity = False
-    shouldCallAfterRemoveLiquidity = False
-    enableHookAdjustedAmounts = False
+    should_call_compute_dynamic_swap_fee = False
+    should_call_before_swap = False
+    should_call_after_swap = False
+    should_call_before_add_liquidity = False
+    should_call_after_add_liquidity = False
+    should_call_before_remove_liquidity = False
+    should_call_after_remove_liquidity = False
+    enable_hook_adjusted_amounts = False
 
     def on_before_add_liquidity(self):
         return False

python/src/hooks/exit_fee_hook.py

@@ -0,0 +1,89 @@
+from src.remove_liquidity import RemoveKind
+from src.maths import mul_down_fixed
+
+
+# This hook implements the ExitFeeHookExample found in mono-repo: https://github.com/balancer/balancer-v3-monorepo/blob/c848c849cb44dc35f05d15858e4fba9f17e92d5e/pkg/pool-hooks/contracts/ExitFeeHookExample.sol
+class ExitFeeHook:
+    def __init__(self):
+        self.should_call_compute_dynamic_swap_fee = False
+        self.should_call_before_swap = False
+        self.should_call_after_swap = False
+        self.should_call_before_add_liquidity = False
+        self.should_call_after_add_liquidity = False
+        self.should_call_before_remove_liquidity = False
+        self.should_call_after_remove_liquidity = True
+        self.enable_hook_adjusted_amounts = True
+
+    def on_before_add_liquidity(self):
+        return {"success": False, "hook_adjusted_balances_scaled18": []}
+
+    def on_after_add_liquidity(
+        self,
+    ):
+        return {"success": False, "hook_adjusted_amounts_in_raw": []}
+
+    def on_before_remove_liquidity(self):
+        return {"success": False, "hook_adjusted_balances_scaled18": []}
+
+    def on_after_remove_liquidity(
+        self,
+        kind,
+        _bpt_amount_in,
+        _amounts_out_scaled18,
+        amounts_out_raw,
+        _balances_scaled18,
+        hook_state,
+    ):
+        if not (
+            isinstance(hook_state, dict)
+            and hook_state is not None
+            and "removeLiquidityHookFeePercentage" in hook_state
+            and "tokens" in hook_state
+        ):
+            raise ValueError("Unexpected hookState")
+
+        # // Our current architecture only supports fees on tokens. Since we must always respect exact `amountsOut`, and
+        # // non-proportional remove liquidity operations would require taking fees in BPT, we only support proportional
+        # // removeLiquidity.
+        if kind != RemoveKind.PROPORTIONAL.value:
+            raise ValueError("ExitFeeHook: Unsupported RemoveKind: ", kind)
+
+        accrued_fees = [0] * len(hook_state["tokens"])
+        hook_adjusted_amounts_out_raw = amounts_out_raw[:]
+        if hook_state["removeLiquidityHookFeePercentage"] > 0:
+            # Charge fees proportional to amounts out of each token
+
+            for i in range(len(amounts_out_raw)):
+                hook_fee = mul_down_fixed(
+                    amounts_out_raw[i],
+                    hook_state["removeLiquidityHookFeePercentage"],
+                )
+                accrued_fees[i] = hook_fee
+                hook_adjusted_amounts_out_raw[i] -= hook_fee
+                # Fees don't need to be transferred to the hook, because donation will reinsert them in the vault
+
+            # // In SC Hook Donates accrued fees back to LPs
+            # // _vault.addLiquidity(
+            # //     AddLiquidityParams({
+            # //         pool: pool,
+            # //         to: msg.sender, // It would mint BPTs to router, but it's a donation so no BPT is minted
+            # //         maxAmountsIn: accruedFees, // Donate all accrued fees back to the pool (i.e. to the LPs)
+            # //         minBptAmountOut: 0, // Donation does not return BPTs, any number above 0 will revert
+            # //         kind: AddLiquidityKind.DONATION,
+            # //         userData: bytes(''), // User data is not used by donation, so we can set to an empty string
+            # //     }),
+            # // );
+
+        return {
+            "success": True,
+            "hook_adjusted_amounts_out_raw": hook_adjusted_amounts_out_raw,
+        }
+
+    def on_before_swap(self):
+        return {"success": False, "hook_adjusted_balances_scaled18": []}
+
+    def on_after_swap(self):
+        return {"success": False, "hook_adjusted_amount_calculated_raw": 0}
+
+    def on_compute_dynamic_swap_fee(self):
+        return {"success": False, "dynamic_swap_fee": 0}

python/src/remove_liquidity.py

@@ -36,13 +36,13 @@ def remove_liquidity(
     )
 
     updated_balances_live_scaled18 = pool_state["balancesLiveScaled18"][:]
-    if hook_class.shouldCallBeforeRemoveLiquidity:
+    if hook_class.should_call_before_remove_liquidity:
         # Note - in SC balances and amounts are updated to reflect any rate change.
         # Daniel said we should not worry about this as any large rate changes
         # will mean something has gone wrong.
         # We do take into account and balance changes due
         # to hook using hookAdjustedBalancesScaled18.
-        hook_return = hook_class.onBeforeRemoveLiquidity(
+        hook_return = hook_class.on_before_remove_liquidity(
             remove_liquidity_input["kind"],
             remove_liquidity_input["max_bpt_amount_in_raw"],
             remove_liquidity_input["min_amounts_out_raw"],
@@ -52,7 +52,7 @@ def remove_liquidity(
         if hook_return["success"] is False:
             raise SystemError("BeforeRemoveLiquidityHookFailed")
 
-        for i, a in enumerate(hook_return["hookAdjustedBalancesScaled18"]):
+        for i, a in enumerate(hook_return["hook_adjusted_balances_scaled18"]):
             updated_balances_live_scaled18[i] = a
 
     if remove_liquidity_input["kind"] == RemoveKind.PROPORTIONAL.value:
@@ -119,24 +119,22 @@ def remove_liquidity(
             swap_fee_amounts_scaled18[i], pool_state["aggregateSwapFee"]
         )
 
-        updated_balances_live_scaled18[i] -= (
+        updated_balances_live_scaled18[i] = updated_balances_live_scaled18[i] - (
             amounts_out_scaled18[i] + aggregate_swap_fee_amount_scaled18
         )
 
-    if hook_class.shouldCallAfterRemoveLiquidity:
+    if hook_class.should_call_after_remove_liquidity:
         hook_return = hook_class.on_after_remove_liquidity(
-            {
-                "kind": remove_liquidity_input["kind"],
-                "bpt_amount_in": bpt_amount_in,
-                "amountsOutScaled18": amounts_out_scaled18,
-                "amountsOutRaw": amounts_out_raw,
-                "updatedBalancesLiveScaled18": updated_balances_live_scaled18,
-                "hook_state": hook_state,
-            }
+            remove_liquidity_input["kind"],
+            bpt_amount_in,
+            amounts_out_scaled18,
+            amounts_out_raw,
+            updated_balances_live_scaled18,
+            hook_state,
         )
 
         if hook_return["success"] is False or len(
-            hook_return["hookAdjustedAmountsOutRaw"]
+            hook_return["hook_adjusted_amounts_out_raw"]
         ) is not len(amounts_out_raw):
             raise SystemError(
                 "AfterRemoveLiquidityHookFailed",
@@ -145,7 +143,7 @@ def remove_liquidity(
             )
 
         # If hook adjusted amounts is not enabled, ignore amounts returned by the hook
-        if hook_class.enableHookAdjustedAmounts:
+        if hook_class.enable_hook_adjusted_amounts:
             for i, a in enumerate(hook_return["hook_adjusted_amounts_out_raw"]):
                 amounts_out_raw[i] = a
 

python/src/swap.py

@@ -39,20 +39,22 @@ def swap(swap_input, pool_state, pool_class, hook_class, hook_state):
     )
 
     updated_balances_live_scaled18 = pool_state["balancesLiveScaled18"][:]
-    if hook_class.shouldCallBeforeSwap:
+    if hook_class.should_call_before_swap:
         # Note - in SC balances and amounts are updated to reflect any rate change.
         # Daniel said we should not worry about this as any large rate changes
         # will mean something has gone wrong.
         # We do take into account and balance changes due
         # to hook using hookAdjustedBalancesScaled18.
-        hook_return = hook_class.onBeforeSwap({**swap_input, "hook_state": hook_state})
+        hook_return = hook_class.on_before_swap(
+            {**swap_input, "hook_state": hook_state}
+        )
         if hook_return["success"] is False:
             raise SystemError("BeforeSwapHookFailed")
-        for i, a in enumerate(hook_return["hookAdjustedBalancesScaled18"]):
+        for i, a in enumerate(hook_return["hook_adjusted_balances_scaled18"]):
             updated_balances_live_scaled18[i] = a
 
     swap_fee = pool_state["swapFee"]
-    if hook_class.shouldCallComputeDynamicSwapFee:
+    if hook_class.should_call_compute_dynamic_swap_fee:
         hook_return = hook_class.onComputeDynamicSwapFee(
             swap_input,
             pool_state["swapFee"],
@@ -115,7 +117,7 @@ def swap(swap_input, pool_state, pool_class, hook_class, hook_state):
             amount_given_scaled18,
             amount_calculated_scaled18 + aggregate_swap_fee_amount_scaled18,
         )
-        if swap_input["swap_kind"] == SwapKind.GIVENIN
+        if swap_input["swap_kind"] == SwapKind.GIVENIN.value
         else (
             amount_calculated_scaled18 - aggregate_swap_fee_amount_scaled18,
             amount_given_scaled18,
@@ -125,20 +127,20 @@ def swap(swap_input, pool_state, pool_class, hook_class, hook_state):
     updated_balances_live_scaled18[input_index] += balance_in_increment
     updated_balances_live_scaled18[output_index] -= balance_out_decrement
 
-    if hook_class.shouldCallAfterSwap:
-        hook_return = hook_class.onAfterSwap(
+    if hook_class.should_call_after_swap:
+        hook_return = hook_class.on_after_swap(
             {
                 "kind": swap_input["swap_kind"],
                 "token_in": swap_input["token_in"],
                 "token_out": swap_input["token_out"],
                 "amount_in_scaled18": (
                     amount_given_scaled18
-                    if swap_input["swap_kind"] == SwapKind.GIVENIN
+                    if swap_input["swap_kind"] == SwapKind.GIVENIN.value
                     else amount_calculated_scaled18
                 ),
                 "amount_out_scaled18": (
                     amount_calculated_scaled18
-                    if swap_input["swap_kind"] == SwapKind.GIVENIN
+                    if swap_input["swap_kind"] == SwapKind.GIVENIN.value
                     else amount_given_scaled18
                 ),
                 "token_in_balance_scaled18": updated_balances_live_scaled18[
@@ -157,7 +159,7 @@ def swap(swap_input, pool_state, pool_class, hook_class, hook_state):
                 "AfterAddSwapHookFailed", pool_state["poolType"], pool_state["hookType"]
             )
         # If hook adjusted amounts is not enabled, ignore amount returned by the hook
-        if hook_class.enableHookAdjustedAmounts:
+        if hook_class.enable_hook_adjusted_amounts:
             amount_calculated_raw = hook_return["hook_adjusted_amount_calculated_raw"]
 
     return amount_calculated_raw

python/src/vault.py

@@ -1,10 +1,11 @@
-from .swap import swap
-from .add_liquidity import add_liquidity
-from .remove_liquidity import remove_liquidity
-from .pools.weighted import Weighted
-from .pools.buffer.erc4626_buffer_wrap_or_unwrap import erc4626_buffer_wrap_or_unwrap
-from .pools.stable import Stable
-from .hooks.default_hook import DefaultHook
+from src.swap import swap
+from src.add_liquidity import add_liquidity
+from src.remove_liquidity import remove_liquidity
+from src.pools.weighted import Weighted
+from src.pools.buffer.erc4626_buffer_wrap_or_unwrap import erc4626_buffer_wrap_or_unwrap
+from src.pools.stable import Stable
+from src.hooks.default_hook import DefaultHook
+from src.hooks.exit_fee_hook import ExitFeeHook
 
 
 class Vault:
@@ -13,10 +14,10 @@ def __init__(self, *, custom_pool_classes=None, custom_hook_classes=None):
             "Weighted": Weighted,
             "Stable": Stable,
         }
+        self.hook_classes = {"ExitFee": ExitFeeHook}
         if custom_pool_classes is not None:
             self.pool_classes.update(custom_pool_classes)
 
-        self.hook_classes = {}
         if custom_hook_classes is not None:
             self.hook_classes.update(custom_hook_classes)
 
@@ -63,4 +64,4 @@ def _get_hook(self, hook_name, hook_state):
             raise SystemError("Unsupported Hook Type:", hook_name)
         if hook_state is None:
             raise SystemError("No state for Hook:", hook_name)
-        return hook_class(hook_state)
+        return hook_class()