Fixed PoolModifyPositionTest and some more breaking tests.

Uniswap · Dec 7, 2023 · 7ec6d8d · 7ec6d8d
1 parent 3fd4079
commit 7ec6d8d
Show file tree

Hide file tree

Showing 31 changed files with 499 additions and 125 deletions.
diff --git a/.forge-snapshots/addLiquidity with empty hook.snap b/.forge-snapshots/addLiquidity with empty hook.snap
@@ -1 +1 @@
-319666
+320115
diff --git a/.forge-snapshots/addLiquidity with native token.snap b/.forge-snapshots/addLiquidity with native token.snap
@@ -1 +1 @@
-202489
+202901
diff --git a/.forge-snapshots/addLiquidity.snap b/.forge-snapshots/addLiquidity.snap
@@ -1 +1 @@
-202409
+202847
diff --git a/.forge-snapshots/before swap hook, already cached dynamic fee.snap b/.forge-snapshots/before swap hook, already cached dynamic fee.snap
@@ -1 +1 @@
-193153
+193154
diff --git a/.forge-snapshots/cached dynamic fee, no hooks.snap b/.forge-snapshots/cached dynamic fee, no hooks.snap
@@ -1 +1 @@
-148516
+148495
diff --git a/.forge-snapshots/donate gas with 1 token.snap b/.forge-snapshots/donate gas with 1 token.snap
@@ -1 +1 @@
-139130
+139131
diff --git a/.forge-snapshots/donate gas with 2 tokens.snap b/.forge-snapshots/donate gas with 2 tokens.snap
@@ -1 +1 @@
-186552
+186576
diff --git a/.forge-snapshots/initialize.snap b/.forge-snapshots/initialize.snap
@@ -1 +1 @@
-78675
+78673
diff --git a/.forge-snapshots/mintWithEmptyHookEOAInitiated.snap b/.forge-snapshots/mintWithEmptyHookEOAInitiated.snap
@@ -1 +1 @@
-250947
+251882
diff --git a/.forge-snapshots/modify position with noop.snap b/.forge-snapshots/modify position with noop.snap
@@ -1 +1 @@
-56515
+56828
diff --git a/.forge-snapshots/poolManager bytecode size.snap b/.forge-snapshots/poolManager bytecode size.snap
@@ -1 +1 @@
-25626
+26246
diff --git a/.forge-snapshots/removeLiquidity with empty hook.snap b/.forge-snapshots/removeLiquidity with empty hook.snap
@@ -0,0 +1 @@
+105888
diff --git a/.forge-snapshots/removeLiquidity with native token.snap b/.forge-snapshots/removeLiquidity with native token.snap
@@ -0,0 +1 @@
+211271
diff --git a/.forge-snapshots/removeLiquidity.snap b/.forge-snapshots/removeLiquidity.snap
@@ -0,0 +1 @@
+207538
diff --git a/.forge-snapshots/simple swap with native.snap b/.forge-snapshots/simple swap with native.snap
@@ -1 +1 @@
-197401
+197380
diff --git a/.forge-snapshots/simple swap.snap b/.forge-snapshots/simple swap.snap
@@ -1 +1 @@
-205969
+205948
diff --git a/.forge-snapshots/simpleSwapEOAInitiated.snap b/.forge-snapshots/simpleSwapEOAInitiated.snap
@@ -1 +1 @@
-177213
+177214
diff --git a/.forge-snapshots/simpleSwapNativeEOAInitiated.snap b/.forge-snapshots/simpleSwapNativeEOAInitiated.snap
@@ -1 +1 @@
-173875
+173876
diff --git a/.forge-snapshots/swap against liquidity with native token.snap b/.forge-snapshots/swap against liquidity with native token.snap
@@ -1 +1 @@
-127994
+127973
diff --git a/.forge-snapshots/swap against liquidity.snap b/.forge-snapshots/swap against liquidity.snap
@@ -1 +1 @@
-115485
+115464
diff --git a/.forge-snapshots/swap burn claim for input.snap b/.forge-snapshots/swap burn claim for input.snap
@@ -1 +1 @@
-134557
+134513
diff --git a/.forge-snapshots/swap mint output as claim.snap b/.forge-snapshots/swap mint output as claim.snap
@@ -1 +1 @@
-218140
+218141
diff --git a/.forge-snapshots/swap with dynamic fee.snap b/.forge-snapshots/swap with dynamic fee.snap
@@ -1 +1 @@
-192395
+192396
diff --git a/.forge-snapshots/swap with hooks.snap b/.forge-snapshots/swap with hooks.snap
@@ -1 +1 @@
-115463
+115442
diff --git a/.forge-snapshots/update dynamic fee in before swap.snap b/.forge-snapshots/update dynamic fee in before swap.snap
@@ -1 +1 @@
-198966
+198967
diff --git a/src/libraries/Hooks.sol b/src/libraries/Hooks.sol
@@ -7,7 +7,7 @@ import {FeeLibrary} from "../libraries/FeeLibrary.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.
+/// has leading bits '1001' which would cause the 'before initialize' and 'after add liquidity' hooks to be used.
 library Hooks {
     using FeeLibrary for uint24;
 

diff --git a/src/test/AccessLockHook.sol b/src/test/AccessLockHook.sol
@@ -76,6 +76,15 @@ contract AccessLockHook is Test, BaseTestHooks {
         return _executeAction(key, hookData, IHooks.beforeAddLiquidity.selector);
     }
 
+    function beforeRemoveLiquidity(
+        address, /* sender **/
+        PoolKey calldata key,
+        IPoolManager.ModifyPositionParams calldata, /* params **/
+        bytes calldata hookData
+    ) external override returns (bytes4) {
+        return _executeAction(key, hookData, IHooks.beforeRemoveLiquidity.selector);
+    }
+
     function _executeAction(PoolKey memory key, bytes calldata hookData, bytes4 selector) internal returns (bytes4) {
         if (hookData.length == 0) {
             // We have re-entered the hook or we are initializing liquidity in the pool before testing the lock actions.

diff --git a/src/test/PoolModifyPositionTest.sol b/src/test/PoolModifyPositionTest.sol
@@ -16,37 +16,46 @@ contract PoolModifyPositionTest is Test, PoolTestBase {
     using Hooks for IHooks;
     using FeeLibrary for uint24;
 
+    enum LockAction {
+        AddLiquidity,
+        RemoveLiquidity
+    }
+
     constructor(IPoolManager _manager) PoolTestBase(_manager) {}
 
     struct CallbackData {
         address sender;
         PoolKey key;
         IPoolManager.ModifyPositionParams params;
         bytes hookData;
+        LockAction action;
     }
 
     function removeLiquidity(PoolKey memory key, IPoolManager.ModifyPositionParams memory params, bytes memory hookData)
         external
         payable
         returns (BalanceDelta delta)
     {
-        delta = _modifyPosition(key, params, hookData);
+        delta = _modifyPosition(key, params, hookData, LockAction.RemoveLiquidity);
     }
 
     function addLiquidity(PoolKey memory key, IPoolManager.ModifyPositionParams memory params, bytes memory hookData)
         external
         payable
         returns (BalanceDelta delta)
     {
-        delta = _modifyPosition(key, params, hookData);
+        delta = _modifyPosition(key, params, hookData, LockAction.AddLiquidity);
     }
 
-    function _modifyPosition(PoolKey memory key, IPoolManager.ModifyPositionParams memory params, bytes memory hookData)
-        internal
-        returns (BalanceDelta delta)
-    {
+    function _modifyPosition(
+        PoolKey memory key,
+        IPoolManager.ModifyPositionParams memory params,
+        bytes memory hookData,
+        LockAction action
+    ) internal returns (BalanceDelta delta) {
         delta = abi.decode(
-            manager.lock(address(this), abi.encode(CallbackData(msg.sender, key, params, hookData))), (BalanceDelta)
+            manager.lock(address(this), abi.encode(CallbackData(msg.sender, key, params, hookData, action))),
+            (BalanceDelta)
         );
 
         uint256 ethBalance = address(this).balance;
@@ -61,7 +70,7 @@ contract PoolModifyPositionTest is Test, PoolTestBase {
         CallbackData memory data = abi.decode(rawData, (CallbackData));
 
         BalanceDelta delta;
-        if (data.params.liquidityDelta >= 0) {
+        if (data.action == LockAction.AddLiquidity) {
             delta = manager.addLiquidity(data.key, data.params, data.hookData);
         } else {
             delta = manager.removeLiquidity(data.key, data.params, data.hookData);
@@ -75,7 +84,7 @@ contract PoolModifyPositionTest is Test, PoolTestBase {
 
         // These assertions only apply in non lock-accessing pools.
         if (!data.key.hooks.hasPermissionToAccessLock()) {
-            if (data.params.liquidityDelta > 0) {
+            if (data.action == LockAction.AddLiquidity) {
                 require(delta0 > 0 || delta1 > 0 || data.key.hooks.hasPermissionToNoOp(), "assert 1 failed");
                 require(!(delta0 < 0 || delta1 < 0), "assert 2 failed");
             } else {