Add StableSurge hook support (python)

python/pyproject.toml

@@ -0,0 +1,4 @@
+[tool.pytest.ini_options]
+pythonpath = [
+  ".", "src"
+]

python/src/hooks/stable_surge_hook.py

@@ -0,0 +1,119 @@
+from typing import Dict, List
+
+from src.maths import (
+    div_down_fixed,
+    mul_down_fixed,
+    complement_fixed,
+)
+from src.pools.stable import Stable
+from src.swap import SwapKind
+
+
+# This hook implements the StableSurgeHook found in mono-repo: https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/pool-hooks/contracts/StableSurgeHook.sol
+class StableSurgeHook:
+    def __init__(self):
+        self.should_call_compute_dynamic_swap_fee = True
+        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 = False
+        self.enable_hook_adjusted_amounts = False
+
+    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):
+        return {"success": False, "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,
+        params: Dict,
+        static_swap_fee_percentage: int,
+        hook_state: Dict,
+    ) -> Dict[str, int]:
+        stable_pool = Stable(hook_state)
+
+        return {
+            "success": True,
+            "dynamic_swap_fee": self.get_surge_fee_percentage(
+                params,
+                stable_pool,
+                hook_state["surgeThresholdPercentage"],
+                hook_state["maxSurgeFeePercentage"],
+                static_swap_fee_percentage,
+            ),
+        }
+
+    def get_surge_fee_percentage(
+        self,
+        params: Dict,
+        pool: Stable,
+        surge_threshold_percentage: int,
+        max_surge_fee_percentage: int,
+        static_fee_percentage: int,
+    ) -> int:
+        amount_calculated_scaled_18 = pool.on_swap(params)
+        new_balances = params["balances_live_scaled18"][:]
+
+        if params["swap_kind"] == SwapKind.GIVENIN.value:
+            new_balances[params["index_in"]] += params["amount_given_scaled18"]
+            new_balances[params["index_out"]] -= amount_calculated_scaled_18
+        else:
+            new_balances[params["index_in"]] += amount_calculated_scaled_18
+            new_balances[params["index_out"]] -= params["amount_given_scaled18"]
+
+        new_total_imbalance = self.calculate_imbalance(new_balances)
+
+        if new_total_imbalance == 0:
+            return static_fee_percentage
+
+        old_total_imbalance = self.calculate_imbalance(params["balances_live_scaled18"])
+
+        if (
+            new_total_imbalance <= old_total_imbalance
+            or new_total_imbalance <= surge_threshold_percentage
+        ):
+            return static_fee_percentage
+
+        dynamic_swap_fee = static_fee_percentage + mul_down_fixed(
+            max_surge_fee_percentage - static_fee_percentage,
+            div_down_fixed(
+                new_total_imbalance - surge_threshold_percentage,
+                complement_fixed(surge_threshold_percentage),
+            ),
+        )
+        return dynamic_swap_fee
+
+    def calculate_imbalance(self, balances: List[int]) -> int:
+        median = self.find_median(balances)
+
+        total_balance = sum(balances)
+        total_diff = sum(self.abs_sub(balance, median) for balance in balances)
+
+        return div_down_fixed(total_diff, total_balance)
+
+    def find_median(self, balances: List[int]) -> int:
+        sorted_balances = sorted(balances)
+        mid = len(sorted_balances) // 2
+
+        if len(sorted_balances) % 2 == 0:
+            return (sorted_balances[mid - 1] + sorted_balances[mid]) // 2
+        else:
+            return sorted_balances[mid]
+
+    def abs_sub(self, a: int, b: int) -> int:
+        return abs(a - b)

python/src/swap.py

@@ -55,16 +55,6 @@ def swap(swap_input, pool_state, pool_class, hook_class, hook_state):
         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.should_call_compute_dynamic_swap_fee:
-        hook_return = hook_class.onComputeDynamicSwapFee(
-            swap_input,
-            pool_state["swapFee"],
-            hook_state,
-        )
-        if hook_return["success"] is True:
-            swap_fee = hook_return["dynamicSwapFee"]
-
     # _swap()
     swap_params = {
         "swap_kind": swap_input["swap_kind"],
@@ -74,6 +64,16 @@ def swap(swap_input, pool_state, pool_class, hook_class, hook_state):
         "index_out": output_index,
     }
 
+    swap_fee = pool_state["swapFee"]
+    if hook_class.should_call_compute_dynamic_swap_fee:
+        hook_return = hook_class.on_compute_dynamic_swap_fee(
+            swap_params,
+            pool_state["swapFee"],
+            hook_state,
+        )
+        if hook_return["success"] is True:
+            swap_fee = hook_return["dynamic_swap_fee"]
+
     total_swap_fee_amount_scaled18 = 0
     if swap_params["swap_kind"] == SwapKind.GIVENIN.value:
         # Round up to avoid losses during precision loss.

python/src/vault.py

@@ -6,6 +6,7 @@
 from src.pools.stable import Stable
 from src.hooks.default_hook import DefaultHook
 from src.hooks.exit_fee_hook import ExitFeeHook
+from src.hooks.stable_surge_hook import StableSurgeHook
 
 
 class Vault:
@@ -14,7 +15,7 @@ def __init__(self, *, custom_pool_classes=None, custom_hook_classes=None):
             "WEIGHTED": Weighted,
             "STABLE": Stable,
         }
-        self.hook_classes = {"ExitFee": ExitFeeHook}
+        self.hook_classes = {"ExitFee": ExitFeeHook, "StableSurge": StableSurgeHook}
         if custom_pool_classes is not None:
             self.pool_classes.update(custom_pool_classes)
 

python/test/hooks/test_stable_surge.py

@@ -0,0 +1,68 @@
+from src.vault import Vault
+from src.swap import SwapKind
+
+pool_state = {
+    "poolType": "STABLE",
+    "hookType": "StableSurge",
+    "poolAddress": "0x132F4bAa39330d9062fC52d81dF72F601DF8C01f",
+    "tokens": [
+        "0x7b79995e5f793a07bc00c21412e50ecae098e7f9",
+        "0xb19382073c7a0addbb56ac6af1808fa49e377b75",
+    ],
+    "scalingFactors": [1, 1],
+    "swapFee": 10000000000000000,
+    "aggregateSwapFee": 10000000000000000,
+    "balancesLiveScaled18": [10000000000000000, 10000000000000000000],
+    "tokenRates": [1000000000000000000, 1000000000000000000],
+    "totalSupply": 9079062661965173292,
+    "amp": 1000000,
+    "supportsUnbalancedLiquidity": True,
+}
+
+hook_state = {
+    "hookType": "StableSurge",
+    "surgeThresholdPercentage": 300000000000000000,
+    "maxSurgeFeePercentage": 950000000000000000,
+    "amp": pool_state["amp"],
+}
+
+vault = Vault()
+
+
+def test_below_surge_threshold_static_swap_fee_case1():
+    swap_input = {
+        "swap_kind": SwapKind.GIVENIN.value,
+        "amount_raw": 1000000000000000,
+        "token_in": pool_state["tokens"][0],
+        "token_out": pool_state["tokens"][1],
+    }
+    output_amount = vault.swap(
+        swap_input=swap_input, pool_state=pool_state, hook_state=hook_state
+    )
+    assert output_amount == 78522716365403684
+
+
+def test_below_surge_threshold_static_swap_fee_case2():
+    swap_input = {
+        "swap_kind": SwapKind.GIVENIN.value,
+        "amount_raw": 10000000000000000,
+        "token_in": pool_state["tokens"][0],
+        "token_out": pool_state["tokens"][1],
+    }
+    output_amount = vault.swap(
+        swap_input=swap_input, pool_state=pool_state, hook_state=hook_state
+    )
+    assert output_amount == 452983383563178802
+
+
+def test_above_surge_threshold_uses_surge_fee():
+    swap_input = {
+        "swap_kind": SwapKind.GIVENIN.value,
+        "amount_raw": 8000000000000000000,
+        "token_in": pool_state["tokens"][1],
+        "token_out": pool_state["tokens"][0],
+    }
+    output_amount = vault.swap(
+        swap_input=swap_input, pool_state=pool_state, hook_state=hook_state
+    )
+    assert output_amount == 3252130027531260