feat: hooks refactor POC

.forge-snapshots/HooksShouldCallBeforeSwap.snap

@@ -1 +1 @@
-34
+116

.forge-snapshots/before swap hook, already cached dynamic fee.snap

@@ -1 +1 @@
-189798
+190434

.forge-snapshots/cached dynamic fee, no hooks.snap

@@ -1 +1 @@
-145314
+145735

.forge-snapshots/donate gas with 1 token.snap

@@ -1 +1 @@
-134915
+135307

.forge-snapshots/donate gas with 2 tokens.snap

@@ -1 +1 @@
-186636
+187028

.forge-snapshots/gas overhead of no-op lock.snap

@@ -1 +1 @@
-14918
+14910

.forge-snapshots/initialize.snap

@@ -1 +1 @@
-53843
+54228

.forge-snapshots/mint with empty hook.snap

@@ -1 +1 @@
-319184
+319684

.forge-snapshots/mint with native token.snap

@@ -1 +1 @@
-198926
+199307

.forge-snapshots/mint.snap

@@ -1 +1 @@
-202041
+202422

.forge-snapshots/poolManager bytecode size.snap

@@ -1 +1 @@
-24805
+25224

.forge-snapshots/simple swap with native.snap

@@ -1 +1 @@
-190722
+191143

.forge-snapshots/simple swap.snap

@@ -1 +1 @@
-202425
+202846

.forge-snapshots/swap against liquidity with native token.snap

@@ -1 +1 @@
-121318
+121739

.forge-snapshots/swap against liquidity.snap

@@ -1 +1 @@
-109143
+109564

.forge-snapshots/swap burn claim for input.snap

@@ -1 +1 @@
-127882
+128303

.forge-snapshots/swap mint output as claim.snap

@@ -1 +1 @@
-212594
+213015

.forge-snapshots/swap with dynamic fee.snap

@@ -1 +1 @@
-189038
+189521

.forge-snapshots/swap with hooks.snap

@@ -1 +1 @@
-109122
+109543

.forge-snapshots/update dynamic fee in before swap.snap

@@ -1 +1 @@
-195628
+196264

src/PoolManager.sol

@@ -28,6 +28,7 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
     using SafeCast for *;
     using Pool for *;
     using Hooks for IHooks;
+    using Hooks for PoolKey;
     using Position for mapping(bytes32 => Position.Info);
     using CurrencyLibrary for Currency;
     using FeeLibrary for uint24;
@@ -108,27 +109,15 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         if (key.currency0 >= key.currency1) revert CurrenciesInitializedOutOfOrder();
         if (!key.hooks.isValidHookAddress(key.fee)) revert Hooks.HookAddressNotValid(address(key.hooks));
 
-        if (key.hooks.shouldCallBeforeInitialize()) {
-            if (key.hooks.beforeInitialize(msg.sender, key, sqrtPriceX96, hookData) != IHooks.beforeInitialize.selector)
-            {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.beforeInitialize(sqrtPriceX96, hookData);
 
         PoolId id = key.toId();
 
         uint24 swapFee = key.fee.isDynamicFee() ? _fetchDynamicSwapFee(key) : key.fee.getStaticFee();
 
         tick = pools[id].initialize(sqrtPriceX96, _fetchProtocolFees(key), _fetchHookFees(key), swapFee);
 
-        if (key.hooks.shouldCallAfterInitialize()) {
-            if (
-                key.hooks.afterInitialize(msg.sender, key, sqrtPriceX96, tick, hookData)
-                    != IHooks.afterInitialize.selector
-            ) {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.afterInitialize(sqrtPriceX96, tick, hookData);
 
         // On intitalize we emit the key's fee, which tells us all fee settings a pool can have: either a static swap fee or dynamic swap fee and if the hook has enabled swap or withdraw fees.
         emit Initialize(id, key.currency0, key.currency1, key.fee, key.tickSpacing, key.hooks);
@@ -185,14 +174,7 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         IPoolManager.ModifyPositionParams memory params,
         bytes calldata hookData
     ) external override noDelegateCall onlyByLocker returns (BalanceDelta delta) {
-        if (key.hooks.shouldCallBeforeModifyPosition()) {
-            if (
-                key.hooks.beforeModifyPosition(msg.sender, key, params, hookData)
-                    != IHooks.beforeModifyPosition.selector
-            ) {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.beforeModifyPosition(params, hookData);
 
         PoolId id = key.toId();
         Pool.FeeAmounts memory feeAmounts;
@@ -223,14 +205,7 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
             }
         }
 
-        if (key.hooks.shouldCallAfterModifyPosition()) {
-            if (
-                key.hooks.afterModifyPosition(msg.sender, key, params, delta, hookData)
-                    != IHooks.afterModifyPosition.selector
-            ) {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.afterModifyPosition(params, delta, hookData);
 
         emit ModifyPosition(id, msg.sender, params.tickLower, params.tickUpper, params.liquidityDelta);
     }
@@ -243,11 +218,7 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         onlyByLocker
         returns (BalanceDelta delta)
     {
-        if (key.hooks.shouldCallBeforeSwap()) {
-            if (key.hooks.beforeSwap(msg.sender, key, params, hookData) != IHooks.beforeSwap.selector) {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.beforeSwap(params, hookData);
 
         PoolId id = key.toId();
 
@@ -276,11 +247,7 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
             }
         }
 
-        if (key.hooks.shouldCallAfterSwap()) {
-            if (key.hooks.afterSwap(msg.sender, key, params, delta, hookData) != IHooks.afterSwap.selector) {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.afterSwap(params, delta, hookData);
 
         emit Swap(
             id, msg.sender, delta.amount0(), delta.amount1(), state.sqrtPriceX96, state.liquidity, state.tick, swapFee
@@ -295,21 +262,13 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         onlyByLocker
         returns (BalanceDelta delta)
     {
-        if (key.hooks.shouldCallBeforeDonate()) {
-            if (key.hooks.beforeDonate(msg.sender, key, amount0, amount1, hookData) != IHooks.beforeDonate.selector) {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.beforeDonate(amount0, amount1, hookData);
 
         delta = _getPool(key).donate(amount0, amount1);
 
         _accountPoolBalanceDelta(key, delta);
 
-        if (key.hooks.shouldCallAfterDonate()) {
-            if (key.hooks.afterDonate(msg.sender, key, amount0, amount1, hookData) != IHooks.afterDonate.selector) {
-                revert Hooks.InvalidHookResponse();
-            }
-        }
+        key.afterDonate(amount0, amount1, hookData);
     }
 
     /// @inheritdoc IPoolManager

src/libraries/Hooks.sol

@@ -2,14 +2,18 @@
 pragma solidity ^0.8.20;
 
 import {IHooks} from "../interfaces/IHooks.sol";
+import {IPoolManager} from "../interfaces/IPoolManager.sol";
 import {FeeLibrary} from "../libraries/FeeLibrary.sol";
+import {BalanceDelta} from "../types/BalanceDelta.sol";
+import {PoolKey} from "../types/PoolKey.sol";
 
 /// @notice V4 decides whether to invoke specific hooks by inspecting the leading bits of the address that
 /// the hooks contract is deployed to.
 /// For example, a hooks contract deployed to address: 0x9000000000000000000000000000000000000000
 /// has leading bits '1001' which would cause the 'before initialize' and 'after modify position' hooks to be used.
 library Hooks {
     using FeeLibrary for uint24;
+    using Hooks for IHooks;
 
     uint256 internal constant BEFORE_INITIALIZE_FLAG = 1 << 159;
     uint256 internal constant AFTER_INITIALIZE_FLAG = 1 << 158;
@@ -44,12 +48,14 @@ library Hooks {
     /// @dev calls param is memory as the function will be called from constructors
     function validateHookAddress(IHooks self, Calls memory calls) internal pure {
         if (
-            calls.beforeInitialize != shouldCallBeforeInitialize(self)
-                || calls.afterInitialize != shouldCallAfterInitialize(self)
-                || calls.beforeModifyPosition != shouldCallBeforeModifyPosition(self)
-                || calls.afterModifyPosition != shouldCallAfterModifyPosition(self)
-                || calls.beforeSwap != shouldCallBeforeSwap(self) || calls.afterSwap != shouldCallAfterSwap(self)
-                || calls.beforeDonate != shouldCallBeforeDonate(self) || calls.afterDonate != shouldCallAfterDonate(self)
+            calls.beforeInitialize != self.shouldCall(BEFORE_INITIALIZE_FLAG)
+                || calls.afterInitialize != self.shouldCall(AFTER_INITIALIZE_FLAG)
+                || calls.beforeModifyPosition != self.shouldCall(BEFORE_MODIFY_POSITION_FLAG)
+                || calls.afterModifyPosition != self.shouldCall(AFTER_MODIFY_POSITION_FLAG)
+                || calls.beforeSwap != self.shouldCall(BEFORE_SWAP_FLAG)
+                || calls.afterSwap != self.shouldCall(AFTER_SWAP_FLAG)
+                || calls.beforeDonate != self.shouldCall(BEFORE_DONATE_FLAG)
+                || calls.afterDonate != self.shouldCall(AFTER_DONATE_FLAG)
         ) {
             revert HookAddressNotValid(address(self));
         }
@@ -67,35 +73,79 @@ library Hooks {
             );
     }
 
-    function shouldCallBeforeInitialize(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & BEFORE_INITIALIZE_FLAG != 0;
+    function validateHooksResponse(bytes4 selector, bytes4 expectedSelector) internal pure {
+        if (selector != expectedSelector) {
+            revert InvalidHookResponse();
+        }
+    }
+
+    function shouldCall(IHooks self, uint256 call) internal pure returns (bool) {
+        return uint256(uint160(address(self))) & call != 0;
+    }
+
+    function beforeInitialize(PoolKey memory key, uint160 sqrtPriceX96, bytes memory hookData) internal {
+        if (!shouldCall(key.hooks, BEFORE_INITIALIZE_FLAG)) return;
+        validateHooksResponse(
+            key.hooks.beforeInitialize(msg.sender, key, sqrtPriceX96, hookData), IHooks.beforeInitialize.selector
+        );
     }
 
-    function shouldCallAfterInitialize(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & AFTER_INITIALIZE_FLAG != 0;
+    function afterInitialize(PoolKey memory key, uint160 sqrtPriceX96, int24 tick, bytes memory hookData) internal {
+        if (!shouldCall(key.hooks, AFTER_INITIALIZE_FLAG)) return;
+        validateHooksResponse(
+            key.hooks.afterInitialize(msg.sender, key, sqrtPriceX96, tick, hookData), IHooks.afterInitialize.selector
+        );
     }
 
-    function shouldCallBeforeModifyPosition(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & BEFORE_MODIFY_POSITION_FLAG != 0;
+    function beforeModifyPosition(
+        PoolKey memory key,
+        IPoolManager.ModifyPositionParams memory params,
+        bytes memory hookData
+    ) internal {
+        if (!key.hooks.shouldCall(BEFORE_MODIFY_POSITION_FLAG)) return;
+        validateHooksResponse(
+            key.hooks.beforeModifyPosition(msg.sender, key, params, hookData), IHooks.beforeModifyPosition.selector
+        );
     }
 
-    function shouldCallAfterModifyPosition(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & AFTER_MODIFY_POSITION_FLAG != 0;
+    function afterModifyPosition(
+        PoolKey memory key,
+        IPoolManager.ModifyPositionParams memory params,
+        BalanceDelta delta,
+        bytes memory hookData
+    ) internal {
+        if (!key.hooks.shouldCall(AFTER_MODIFY_POSITION_FLAG)) return;
+        validateHooksResponse(
+            key.hooks.afterModifyPosition(msg.sender, key, params, delta, hookData), IHooks.afterModifyPosition.selector
+        );
     }
 
-    function shouldCallBeforeSwap(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & BEFORE_SWAP_FLAG != 0;
+    function beforeSwap(PoolKey memory key, IPoolManager.SwapParams memory params, bytes memory hookData) internal {
+        if (!key.hooks.shouldCall(BEFORE_SWAP_FLAG)) return;
+        validateHooksResponse(key.hooks.beforeSwap(msg.sender, key, params, hookData), IHooks.beforeSwap.selector);
     }
 
-    function shouldCallAfterSwap(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & AFTER_SWAP_FLAG != 0;
+    function afterSwap(
+        PoolKey memory key,
+        IPoolManager.SwapParams memory params,
+        BalanceDelta delta,
+        bytes memory hookData
+    ) internal {
+        if (!key.hooks.shouldCall(AFTER_SWAP_FLAG)) return;
+        validateHooksResponse(key.hooks.afterSwap(msg.sender, key, params, delta, hookData), IHooks.afterSwap.selector);
     }
 
-    function shouldCallBeforeDonate(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & BEFORE_DONATE_FLAG != 0;
+    function beforeDonate(PoolKey memory key, uint256 amount0, uint256 amount1, bytes memory hookData) internal {
+        if (!key.hooks.shouldCall(BEFORE_DONATE_FLAG)) return;
+        validateHooksResponse(
+            key.hooks.beforeDonate(msg.sender, key, amount0, amount1, hookData), IHooks.beforeDonate.selector
+        );
     }
 
-    function shouldCallAfterDonate(IHooks self) internal pure returns (bool) {
-        return uint256(uint160(address(self))) & AFTER_DONATE_FLAG != 0;
+    function afterDonate(PoolKey memory key, uint256 amount0, uint256 amount1, bytes memory hookData) internal {
+        if (!key.hooks.shouldCall(AFTER_DONATE_FLAG)) return;
+        validateHooksResponse(
+            key.hooks.afterDonate(msg.sender, key, amount0, amount1, hookData), IHooks.afterDonate.selector
+        );
     }
 }