Introduce standard batched call under a lock

.forge-snapshots/FullRangeAddInitialLiquidity.snap

@@ -1 +1 @@
-311073
+311137

.forge-snapshots/FullRangeAddLiquidity.snap

@@ -1 +1 @@
-122882
+122946

.forge-snapshots/FullRangeFirstSwap.snap

@@ -1 +1 @@
-80283
+80287

.forge-snapshots/FullRangeInitialize.snap

@@ -1 +1 @@
-1015169
+1015181

.forge-snapshots/FullRangeRemoveLiquidity.snap

@@ -1 +1 @@
-110476
+110544

.forge-snapshots/FullRangeRemoveLiquidityAndRebalance.snap

@@ -1 +1 @@
-239954
+240022

.forge-snapshots/FullRangeSecondSwap.snap

@@ -1 +1 @@
-45993
+45997

.forge-snapshots/FullRangeSwap.snap

@@ -1 +1 @@
-79414
+79418

.forge-snapshots/TWAMMSubmitOrder.snap

@@ -1 +1 @@
-122355
+122359

README.md

@@ -46,7 +46,7 @@ contract CoolHook is BaseHook {
         address,
         IPoolManager.PoolKey calldata key,
         IPoolManager.ModifyLiquidityParams calldata params
-    ) external override poolManagerOnly returns (bytes4) {
+    ) external override onlyByManager returns (bytes4) {
         // hook logic
         return BaseHook.beforeAddLiquidity.selector;
     }

contracts/BaseHook.sol

@@ -6,29 +6,20 @@ import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
 import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {SafeCallback} from "./base/SafeCallback.sol";
+import {ImmutableState} from "./base/ImmutableState.sol";
 import {BeforeSwapDelta} from "@uniswap/v4-core/src/types/BeforeSwapDelta.sol";
 
-abstract contract BaseHook is IHooks {
-    error NotPoolManager();
+abstract contract BaseHook is IHooks, SafeCallback {
     error NotSelf();
     error InvalidPool();
     error LockFailure();
     error HookNotImplemented();
 
-    /// @notice The address of the pool manager
-    IPoolManager public immutable poolManager;
-
-    constructor(IPoolManager _poolManager) {
-        poolManager = _poolManager;
+    constructor(IPoolManager _poolManager) ImmutableState(_poolManager) {
         validateHookAddress(this);
     }
 
-    /// @dev Only the pool manager may call this function
-    modifier poolManagerOnly() {
-        if (msg.sender != address(poolManager)) revert NotPoolManager();
-        _;
-    }
-
     /// @dev Only this address may call this function
     modifier selfOnly() {
         if (msg.sender != address(this)) revert NotSelf();
@@ -50,7 +41,7 @@ abstract contract BaseHook is IHooks {
         Hooks.validateHookPermissions(_this, getHookPermissions());
     }
 
-    function unlockCallback(bytes calldata data) external virtual poolManagerOnly returns (bytes memory) {
+    function _unlockCallback(bytes calldata data) internal virtual override returns (bytes memory) {
         (bool success, bytes memory returnData) = address(this).call(data);
         if (success) return returnData;
         if (returnData.length == 0) revert LockFailure();

contracts/SimpleBatchCall.sol

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {LockAndBatchCall} from "./base/LockAndBatchCall.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {ImmutableState} from "./base/ImmutableState.sol";
+import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol";
+import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+import {TransientStateLibrary} from "@uniswap/v4-core/src/libraries/TransientStateLibrary.sol";
+import {CurrencySettler} from "@uniswap/v4-core/test/utils/CurrencySettler.sol";
+
+/// @title SimpleBatchCall
+/// @notice Implements a naive settle function to perform any arbitrary batch call under one lock to modifyPosition, donate, intitialize, or swap.
+contract SimpleBatchCall is LockAndBatchCall {
+    using CurrencyLibrary for Currency;
+    using TransientStateLibrary for IPoolManager;
+    using CurrencySettler for Currency;
+
+    constructor(IPoolManager _poolManager) ImmutableState(_poolManager) {}
+
+    struct SettleConfig {
+        bool takeClaims;
+        bool settleUsingBurn; // If true, sends the underlying ERC20s.
+    }
+
+    /// @notice We naively settle all currencies that are touched by the batch call. This data is passed in intially to `execute`.
+    function _settle(address sender, bytes memory data) internal override returns (bytes memory settleData) {
+        if (data.length != 0) {
+            (Currency[] memory currenciesTouched, SettleConfig memory config) =
+                abi.decode(data, (Currency[], SettleConfig));
+
+            for (uint256 i = 0; i < currenciesTouched.length; i++) {
+                Currency currency = currenciesTouched[i];
+                int256 delta = poolManager.currencyDelta(address(this), currenciesTouched[i]);
+
+                if (delta < 0) {
+                    currency.settle(poolManager, sender, uint256(-delta), config.settleUsingBurn);
+                }
+                if (delta > 0) {
+                    currency.take(poolManager, address(this), uint256(delta), config.takeClaims);
+                }
+            }
+        }
+    }
+
+    function _handleAfterExecute(bytes memory, /*callReturnData*/ bytes memory /*settleReturnData*/ )
+        internal
+        pure
+        override
+    {
+        return;
+    }
+}

contracts/base/CallsWithLock.sol

@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+pragma solidity ^0.8.19;
+
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {ImmutableState} from "./ImmutableState.sol";
+import {ICallsWithLock} from "../interfaces/ICallsWithLock.sol";
+import {BalanceDelta} from "@uniswap/v4-core/src/types/BalanceDelta.sol";
+
+/// @title CallsWithLock
+/// @notice Handles all the calls to the pool manager contract. Assumes the integrating contract has already acquired a lock.
+abstract contract CallsWithLock is ICallsWithLock, ImmutableState {
+    error NotSelf();
+
+    modifier onlyBySelf() {
+        if (msg.sender != address(this)) revert NotSelf();
+        _;
+    }
+
+    function initializeWithLock(PoolKey memory key, uint160 sqrtPriceX96, bytes calldata hookData)
+        external
+        onlyBySelf
+        returns (bytes memory)
+    {
+        return abi.encode(poolManager.initialize(key, sqrtPriceX96, hookData));
+    }
+
+    function modifyPositionWithLock(
+        PoolKey calldata key,
+        IPoolManager.ModifyLiquidityParams calldata params,
+        bytes calldata hookData
+    ) external onlyBySelf returns (bytes memory) {
+        (BalanceDelta delta, BalanceDelta feeDelta) = poolManager.modifyLiquidity(key, params, hookData);
+        return abi.encode(delta, feeDelta);
+    }
+
+    function swapWithLock(PoolKey memory key, IPoolManager.SwapParams memory params, bytes calldata hookData)
+        external
+        onlyBySelf
+        returns (bytes memory)
+    {
+        return abi.encode(poolManager.swap(key, params, hookData));
+    }
+
+    function donateWithLock(PoolKey memory key, uint256 amount0, uint256 amount1, bytes calldata hookData)
+        external
+        onlyBySelf
+        returns (bytes memory)
+    {
+        return abi.encode(poolManager.donate(key, amount0, amount1, hookData));
+    }
+}

contracts/base/ImmutableState.sol

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+pragma solidity ^0.8.19;
+
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+
+contract ImmutableState {
+    IPoolManager public immutable poolManager;
+
+    constructor(IPoolManager _manager) {
+        poolManager = _manager;
+    }
+}

contracts/base/LockAndBatchCall.sol

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+pragma solidity ^0.8.19;
+
+import {SafeCallback} from "./SafeCallback.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
+import {CallsWithLock} from "./CallsWithLock.sol";
+
+abstract contract LockAndBatchCall is CallsWithLock, SafeCallback {
+    error CallFail(bytes reason);
+
+    function _settle(address sender, bytes memory data) internal virtual returns (bytes memory settleData);
+    function _handleAfterExecute(bytes memory callReturnData, bytes memory settleReturnData) internal virtual;
+
+    /// @param executeData The function selectors and calldata for any of the function selectors in ICallsWithLock encoded as an array of bytes.
+    function execute(bytes memory executeData, bytes memory settleData) external {
+        (bytes memory lockReturnData) =
+            poolManager.unlock(abi.encode(executeData, abi.encode(msg.sender, settleData)));
+        (bytes memory executeReturnData, bytes memory settleReturnData) = abi.decode(lockReturnData, (bytes, bytes));
+        _handleAfterExecute(executeReturnData, settleReturnData);
+    }
+
+    /// @param data This data is passed from the top-level execute function to the internal _executeWithLockCalls and _settle function. It is decoded as two separate dynamic bytes parameters.
+    /// @dev _unlockCallback is responsible for executing the internal calls under the lock and settling open deltas left on the pool
+    function _unlockCallback(bytes calldata data) internal override returns (bytes memory) {
+        (bytes memory executeData, bytes memory settleDataWithSender) = abi.decode(data, (bytes, bytes));
+        (address sender, bytes memory settleData) = abi.decode(settleDataWithSender, (address, bytes));
+        return abi.encode(_executeWithLockCalls(executeData), _settle(sender, settleData));
+    }
+
+    function _executeWithLockCalls(bytes memory data) internal returns (bytes memory) {
+        bytes[] memory calls = abi.decode(data, (bytes[]));
+        bytes[] memory callsReturnData = new bytes[](calls.length);
+
+        for (uint256 i = 0; i < calls.length; i++) {
+            (bool success, bytes memory returnData) = address(this).call(calls[i]);
+            if (!success) revert(string(returnData));
+            callsReturnData[i] = returnData;
+        }
+        return abi.encode(callsReturnData);
+    }
+}

contracts/base/SafeCallback.sol

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {IUnlockCallback} from "@uniswap/v4-core/src/interfaces/callback/IUnlockCallback.sol";
+import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
+import {ImmutableState} from "./ImmutableState.sol";
+
+abstract contract SafeCallback is ImmutableState, IUnlockCallback {
+    error NotManager();
+
+    modifier onlyByManager() {
+        if (msg.sender != address(poolManager)) revert NotManager();
+        _;
+    }
+
+    /// @dev There is no way to force the onlyByManager modifier but for this callback to be safe, it MUST check that the msg.sender is the pool manager.
+    function unlockCallback(bytes calldata data) external onlyByManager returns (bytes memory) {
+        return _unlockCallback(data);
+    }
+
+    function _unlockCallback(bytes calldata data) internal virtual returns (bytes memory);
+}

contracts/hooks/examples/FullRange.sol

@@ -26,7 +26,7 @@ import {BeforeSwapDelta, BeforeSwapDeltaLibrary} from "@uniswap/v4-core/src/type
 
 import "../../libraries/LiquidityAmounts.sol";
 
-contract FullRange is BaseHook, IUnlockCallback {
+contract FullRange is BaseHook {
     using CurrencyLibrary for Currency;
     using CurrencySettler for Currency;
     using PoolIdLibrary for PoolKey;
@@ -295,10 +295,9 @@ contract FullRange is BaseHook, IUnlockCallback {
         pool.hasAccruedFees = false;
     }
 
-    function unlockCallback(bytes calldata rawData)
-        external
-        override(IUnlockCallback, BaseHook)
-        poolManagerOnly
+    function _unlockCallback(bytes calldata rawData)
+        internal
+        override
         returns (bytes memory)
     {
         CallbackData memory data = abi.decode(rawData, (CallbackData));

contracts/hooks/examples/GeomeanOracle.sol

@@ -86,7 +86,7 @@ contract GeomeanOracle is BaseHook {
         external
         view
         override
-        poolManagerOnly
+        onlyByManager
         returns (bytes4)
     {
         // This is to limit the fragmentation of pools using this oracle hook. In other words,
@@ -99,7 +99,7 @@ contract GeomeanOracle is BaseHook {
     function afterInitialize(address, PoolKey calldata key, uint160, int24, bytes calldata)
         external
         override
-        poolManagerOnly
+        onlyByManager
         returns (bytes4)
     {
         PoolId id = key.toId();
@@ -124,7 +124,8 @@ contract GeomeanOracle is BaseHook {
         PoolKey calldata key,
         IPoolManager.ModifyLiquidityParams calldata params,
         bytes calldata
-    ) external override poolManagerOnly returns (bytes4) {
+    ) external override onlyByManager returns (bytes4) {
+        if (params.liquidityDelta < 0) revert OraclePoolMustLockLiquidity();
         int24 maxTickSpacing = poolManager.MAX_TICK_SPACING();
         if (
             params.tickLower != TickMath.minUsableTick(maxTickSpacing)
@@ -139,14 +140,14 @@ contract GeomeanOracle is BaseHook {
         PoolKey calldata,
         IPoolManager.ModifyLiquidityParams calldata,
         bytes calldata
-    ) external view override poolManagerOnly returns (bytes4) {
+    ) external view override onlyByManager returns (bytes4) {
         revert OraclePoolMustLockLiquidity();
     }
 
     function beforeSwap(address, PoolKey calldata key, IPoolManager.SwapParams calldata, bytes calldata)
         external
         override
-        poolManagerOnly
+        onlyByManager
         returns (bytes4, BeforeSwapDelta, uint24)
     {
         _updatePool(key);

contracts/hooks/examples/LimitOrder.sol

@@ -129,7 +129,7 @@ contract LimitOrder is BaseHook {
     function afterInitialize(address, PoolKey calldata key, uint160, int24 tick, bytes calldata)
         external
         override
-        poolManagerOnly
+        onlyByManager
         returns (bytes4)
     {
         setTickLowerLast(key.toId(), getTickLower(tick, key.tickSpacing));
@@ -142,7 +142,7 @@ contract LimitOrder is BaseHook {
         IPoolManager.SwapParams calldata params,
         BalanceDelta,
         bytes calldata
-    ) external override poolManagerOnly returns (bytes4, int128) {
+    ) external override onlyByManager returns (bytes4, int128) {
         (int24 tickLower, int24 lower, int24 upper) = _getCrossedTicks(key.toId(), key.tickSpacing);
         if (lower > upper) return (LimitOrder.afterSwap.selector, 0);
 
@@ -197,7 +197,7 @@ contract LimitOrder is BaseHook {
 
     function _unlockCallbackFill(PoolKey calldata key, int24 tickLower, int256 liquidityDelta)
         private
-        poolManagerOnly
+        onlyByManager
         returns (uint128 amount0, uint128 amount1)
     {
         (BalanceDelta delta,) = poolManager.modifyLiquidity(