rebase

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

@@ -1 +1 @@
-189798
+192134

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

@@ -1 +1 @@
-145314
+145580

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

@@ -1 +1 @@
-134915
+137163

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

@@ -1 +1 @@
-186636
+188884

.forge-snapshots/initialize.snap

@@ -1 +1 @@
-53843
+53880

.forge-snapshots/mint with empty hook.snap

@@ -1 +1 @@
-319184
+341339

.forge-snapshots/mint with native token.snap

@@ -1 +1 @@
-198926
+201148

.forge-snapshots/mint.snap

@@ -1 +1 @@
-202041
+204226

.forge-snapshots/poolManager bytecode size.snap

@@ -1 +1 @@
-24805
+25262

.forge-snapshots/simple swap with native.snap

@@ -1 +1 @@
-190722
+192988

.forge-snapshots/simple swap.snap

@@ -1 +1 @@
-202425
+204691

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

@@ -1 +1 @@
-121318
+121584

.forge-snapshots/swap against liquidity.snap

@@ -1 +1 @@
-109143
+109409

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

@@ -1 +1 @@
-127882
+128148

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

@@ -0,0 +1 @@
+169819

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

@@ -1 +1 @@
-212594
+214860

.forge-snapshots/swap with 1155 as input.snap

@@ -0,0 +1 @@
+156249

.forge-snapshots/swap with dynamic fee.snap

@@ -1 +1 @@
-189038
+191374

.forge-snapshots/swap with hooks.snap

@@ -1 +1 @@
-109122
+109388

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

@@ -1 +1 @@
-195628
+197964

src/PoolManager.sol

@@ -22,6 +22,8 @@ import {PoolId, PoolIdLibrary} from "./types/PoolId.sol";
 import {BalanceDelta} from "./types/BalanceDelta.sol";
 import {Lockers} from "./libraries/Lockers.sol";
 
+import "forge-std/console2.sol";
+
 /// @notice Holds the state for all pools
 contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
     using PoolIdLibrary for PoolKey;
@@ -49,6 +51,9 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
 
     constructor(uint256 controllerGasLimit) Fees(controllerGasLimit) {}
 
+    // todo change to transient storage
+    address currentHook;
+
     function _getPool(PoolKey memory key) private view returns (Pool.State storage) {
         return pools[key.toId()];
     }
@@ -174,8 +179,13 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
     }
 
     modifier onlyByLocker() {
+        // todo fix stack too deep :/
         address locker = Lockers.getCurrentLocker();
-        if (msg.sender != locker) revert LockedBy(locker);
+        if (msg.sender != locker) {
+            if (msg.sender != currentHook && !Hooks.shouldAccessLock(IHooks(currentHook))) {
+                revert LockedBy(locker);
+            }
+        }
         _;
     }
 
@@ -185,11 +195,15 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         IPoolManager.ModifyPositionParams memory params,
         bytes calldata hookData
     ) external override noDelegateCall onlyByLocker returns (BalanceDelta delta) {
+        _setCurrentHook(address(key.hooks));
+
         if (key.hooks.shouldCallBeforeModifyPosition()) {
-            if (
-                key.hooks.beforeModifyPosition(msg.sender, key, params, hookData)
-                    != IHooks.beforeModifyPosition.selector
-            ) {
+            bytes4 hookReturn = key.hooks.beforeModifyPosition(msg.sender, key, params, hookData);
+
+            if (hookReturn == Hooks.OVERRIDE_SELECTOR) {
+                return delta;
+            }
+            if (hookReturn != IHooks.beforeModifyPosition.selector) {
                 revert Hooks.InvalidHookResponse();
             }
         }
@@ -243,8 +257,14 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         onlyByLocker
         returns (BalanceDelta delta)
     {
+        _setCurrentHook(address(key.hooks));
+
         if (key.hooks.shouldCallBeforeSwap()) {
-            if (key.hooks.beforeSwap(msg.sender, key, params, hookData) != IHooks.beforeSwap.selector) {
+            bytes4 hookReturn = key.hooks.beforeSwap(msg.sender, key, params, hookData);
+            if (hookReturn == Hooks.OVERRIDE_SELECTOR) {
+                return delta;
+            }
+            if (hookReturn != IHooks.beforeSwap.selector) {
                 revert Hooks.InvalidHookResponse();
             }
         }
@@ -295,8 +315,13 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         onlyByLocker
         returns (BalanceDelta delta)
     {
+        _setCurrentHook(address(key.hooks));
         if (key.hooks.shouldCallBeforeDonate()) {
-            if (key.hooks.beforeDonate(msg.sender, key, amount0, amount1, hookData) != IHooks.beforeDonate.selector) {
+            bytes4 hookReturn = key.hooks.beforeDonate(msg.sender, key, amount0, amount1, hookData);
+            if (hookReturn == Hooks.OVERRIDE_SELECTOR) {
+                return delta;
+            }
+            if (hookReturn != IHooks.beforeDonate.selector) {
                 revert Hooks.InvalidHookResponse();
             }
         }
@@ -393,6 +418,11 @@ contract PoolManager is IPoolManager, Fees, NoDelegateCall, Claims {
         return Lockers.nonzeroDeltaCount();
     }
 
+    // TODO: Use transient storage
+    function _setCurrentHook(address hookAddr) internal {
+        currentHook = hookAddr;
+    }
+
     /// @notice receive native tokens for native pools
     receive() external payable {}
 }

src/libraries/Hooks.sol

@@ -4,6 +4,8 @@ pragma solidity ^0.8.20;
 import {IHooks} from "../interfaces/IHooks.sol";
 import {FeeLibrary} from "../libraries/FeeLibrary.sol";
 
+import "forge-std/console2.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
@@ -19,7 +21,12 @@ library Hooks {
     uint256 internal constant AFTER_SWAP_FLAG = 1 << 154;
     uint256 internal constant BEFORE_DONATE_FLAG = 1 << 153;
     uint256 internal constant AFTER_DONATE_FLAG = 1 << 152;
+    uint256 internal constant ACCESS_LOCK_FLAG = 1 << 151;
+    uint256 internal constant OVERRIDE_FLAG = 1 << 150;
+
+    bytes4 constant OVERRIDE_SELECTOR = bytes4(keccak256("OVERRIDE_SELECTOR"));
 
+    // todo not necessarily all calls now?
     struct Calls {
         bool beforeInitialize;
         bool afterInitialize;
@@ -29,6 +36,8 @@ library Hooks {
         bool afterSwap;
         bool beforeDonate;
         bool afterDonate;
+        bool accessLock;
+        bool overrideSelector;
     }
 
     /// @notice Thrown if the address will not lead to the specified hook calls being called
@@ -50,6 +59,7 @@ library Hooks {
                 || calls.afterModifyPosition != shouldCallAfterModifyPosition(self)
                 || calls.beforeSwap != shouldCallBeforeSwap(self) || calls.afterSwap != shouldCallAfterSwap(self)
                 || calls.beforeDonate != shouldCallBeforeDonate(self) || calls.afterDonate != shouldCallAfterDonate(self)
+                || calls.accessLock != shouldAccessLock(self) || calls.overrideSelector != shouldAllowOverride(self)
         ) {
             revert HookAddressNotValid(address(self));
         }
@@ -62,7 +72,7 @@ library Hooks {
         return address(hook) == address(0)
             ? !fee.isDynamicFee() && !fee.hasHookSwapFee() && !fee.hasHookWithdrawFee()
             : (
-                uint160(address(hook)) >= AFTER_DONATE_FLAG || fee.isDynamicFee() || fee.hasHookSwapFee()
+                uint160(address(hook)) >= ACCESS_LOCK_FLAG || fee.isDynamicFee() || fee.hasHookSwapFee()
                     || fee.hasHookWithdrawFee()
             );
     }
@@ -98,4 +108,12 @@ library Hooks {
     function shouldCallAfterDonate(IHooks self) internal pure returns (bool) {
         return uint256(uint160(address(self))) & AFTER_DONATE_FLAG != 0;
     }
+
+    function shouldAccessLock(IHooks self) internal pure returns (bool) {
+        return uint256(uint160(address(self))) & ACCESS_LOCK_FLAG != 0;
+    }
+
+    function shouldAllowOverride(IHooks self) internal pure returns (bool) {
+        return (uint256(uint160(address(self))) & OVERRIDE_FLAG) != 0;
+    }
 }

src/test/AccessLockHook.sol

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.20;
+
+import {BaseTestHooks} from "./BaseTestHooks.sol";
+import {PoolKey} from "../types/PoolKey.sol";
+import {IPoolManager} from "../interfaces/IPoolManager.sol";
+import {IHooks} from "../interfaces/IHooks.sol";
+import {CurrencyLibrary, Currency} from "../types/Currency.sol";
+import {Hooks} from "../libraries/Hooks.sol";
+
+import "forge-std/console2.sol";
+
+contract AccessLockHook is BaseTestHooks {
+    using CurrencyLibrary for Currency;
+
+    IPoolManager manager;
+
+    constructor(IPoolManager _manager) {
+        manager = _manager;
+    }
+
+    function beforeModifyPosition(
+        address sender,
+        PoolKey calldata key,
+        IPoolManager.ModifyPositionParams calldata params,
+        bytes calldata hookData
+    ) external override returns (bytes4) {
+        // just deal with positive deposits of currency1
+        (uint256 amount1) = abi.decode(hookData, (uint256));
+        manager.mint(key.currency1, address(this), amount1);
+        return Hooks.OVERRIDE_SELECTOR;
+    }
+
+    function onERC1155Received(address, address, uint256, uint256, bytes memory) public returns (bytes4) {
+        return this.onERC1155Received.selector;
+    }
+}

src/test/EmptyTestHooks.sol

@@ -20,7 +20,9 @@ contract EmptyTestHooks is IHooks {
                 beforeSwap: true,
                 afterSwap: true,
                 beforeDonate: true,
-                afterDonate: true
+                afterDonate: true,
+                accessLock: false,
+                overrideSelector: false
             })
         );
     }

src/test/PoolModifyPositionTest.sol

@@ -6,6 +6,9 @@ import {IPoolManager} from "../interfaces/IPoolManager.sol";
 import {BalanceDelta} from "../types/BalanceDelta.sol";
 import {PoolKey} from "../types/PoolKey.sol";
 import {PoolTestBase} from "./PoolTestBase.sol";
+import {Hooks} from "../libraries/Hooks.sol";
+
+import "forge-std/console2.sol";
 
 contract PoolModifyPositionTest is PoolTestBase {
     using CurrencyLibrary for Currency;
@@ -45,13 +48,13 @@ contract PoolModifyPositionTest is PoolTestBase {
         if (data.params.liquidityDelta > 0) {
             assert(delta0 > 0 || delta1 > 0);
             assert(!(delta0 < 0 || delta1 < 0));
-            if (delta0 > 0) _settle(data.key.currency0, data.sender, delta.amount0(), true);
-            if (delta1 > 0) _settle(data.key.currency1, data.sender, delta.amount1(), true);
+            if (delta0 > 0) _settle(data.key.currency0, data.sender, int128(delta0), true);
+            if (delta1 > 0) _settle(data.key.currency1, data.sender, int128(delta1), true);
         } else {
             assert(delta0 < 0 || delta1 < 0);
             assert(!(delta0 > 0 || delta1 > 0));
-            if (delta0 < 0) _take(data.key.currency0, data.sender, delta.amount0(), true);
-            if (delta1 < 0) _take(data.key.currency1, data.sender, delta.amount1(), true);
+            if (delta0 < 0) _take(data.key.currency0, data.sender, int128(delta0), true);
+            if (delta1 < 0) _take(data.key.currency1, data.sender, int128(delta1), true);
         }
 
         return abi.encode(delta);