just have 1 check function (#714)

* just have 1 check function * remove double pool lookup --------- Co-authored-by: Alice Henshaw <[email protected]>
Uniswap · May 24, 2024 · 3118115 · 3118115
1 parent 05f986d
commit 3118115
Show file tree

Hide file tree

Showing 34 changed files with 51 additions and 52 deletions.
diff --git a/.forge-snapshots/add liquidity to already existing position with salt.snap b/.forge-snapshots/add liquidity to already existing position with salt.snap
@@ -1 +1 @@
-149414
+149347
diff --git a/.forge-snapshots/addLiquidity CA fee.snap b/.forge-snapshots/addLiquidity CA fee.snap
@@ -1 +1 @@
-326118
+326051
diff --git a/.forge-snapshots/addLiquidity with empty hook.snap b/.forge-snapshots/addLiquidity with empty hook.snap
@@ -1 +1 @@
-280232
+280165
diff --git a/.forge-snapshots/addLiquidity with native token.snap b/.forge-snapshots/addLiquidity with native token.snap
@@ -1 +1 @@
-139620
+139553
diff --git a/.forge-snapshots/create new liquidity to a position with salt.snap b/.forge-snapshots/create new liquidity to a position with salt.snap
@@ -1 +1 @@
-297592
+297525
diff --git a/.forge-snapshots/donate gas with 1 token.snap b/.forge-snapshots/donate gas with 1 token.snap
@@ -1 +1 @@
-106629
+106609
diff --git a/.forge-snapshots/donate gas with 2 tokens.snap b/.forge-snapshots/donate gas with 2 tokens.snap
@@ -1 +1 @@
-147124
+147104
diff --git a/.forge-snapshots/poolManager bytecode size.snap b/.forge-snapshots/poolManager bytecode size.snap
@@ -1 +1 @@
-20003
+19971
diff --git a/.forge-snapshots/removeLiquidity CA fee.snap b/.forge-snapshots/removeLiquidity CA fee.snap
@@ -1 +1 @@
-181711
+181644
diff --git a/.forge-snapshots/removeLiquidity with empty hook.snap b/.forge-snapshots/removeLiquidity with empty hook.snap
@@ -1 +1 @@
-136000
+135933
diff --git a/.forge-snapshots/removeLiquidity with native token.snap b/.forge-snapshots/removeLiquidity with native token.snap
@@ -1 +1 @@
-115828
+115761
diff --git a/.forge-snapshots/set protocol fee.snap b/.forge-snapshots/set protocol fee.snap
@@ -1 +1 @@
-32444
+32478
diff --git a/.forge-snapshots/simple addLiquidity second addition same range.snap b/.forge-snapshots/simple addLiquidity second addition same range.snap
@@ -1 +1 @@
-102856
+102789
diff --git a/.forge-snapshots/simple addLiquidity.snap b/.forge-snapshots/simple addLiquidity.snap
@@ -1 +1 @@
-165348
+165281
diff --git a/.forge-snapshots/simple removeLiquidity some liquidity remains.snap b/.forge-snapshots/simple removeLiquidity some liquidity remains.snap
@@ -1 +1 @@
-96571
+96504
diff --git a/.forge-snapshots/simple removeLiquidity.snap b/.forge-snapshots/simple removeLiquidity.snap
@@ -1 +1 @@
-88611
+88544
diff --git a/.forge-snapshots/simple swap with native.snap b/.forge-snapshots/simple swap with native.snap
@@ -1 +1 @@
-114950
+114887
diff --git a/.forge-snapshots/simple swap.snap b/.forge-snapshots/simple swap.snap
@@ -1 +1 @@
-130115
+130052
diff --git a/.forge-snapshots/swap CA custom curve + swap noop.snap b/.forge-snapshots/swap CA custom curve + swap noop.snap
@@ -1 +1 @@
-132084
+132021
diff --git a/.forge-snapshots/swap CA fee on unspecified.snap b/.forge-snapshots/swap CA fee on unspecified.snap
@@ -1 +1 @@
-179590
+179527
diff --git a/.forge-snapshots/swap against liquidity with native token.snap b/.forge-snapshots/swap against liquidity with native token.snap
@@ -1 +1 @@
-110342
+110279
diff --git a/.forge-snapshots/swap against liquidity.snap b/.forge-snapshots/swap against liquidity.snap
@@ -1 +1 @@
-121689
+121626
diff --git a/.forge-snapshots/swap burn 6909 for input.snap b/.forge-snapshots/swap burn 6909 for input.snap
@@ -1 +1 @@
-133518
+133455
diff --git a/.forge-snapshots/swap burn native 6909 for input.snap b/.forge-snapshots/swap burn native 6909 for input.snap
@@ -1 +1 @@
-122806
+122743
diff --git a/.forge-snapshots/swap mint native output as 6909.snap b/.forge-snapshots/swap mint native output as 6909.snap
@@ -1 +1 @@
-144594
+144531
diff --git a/.forge-snapshots/swap mint output as 6909.snap b/.forge-snapshots/swap mint output as 6909.snap
@@ -1 +1 @@
-161224
+161161
diff --git a/.forge-snapshots/swap skips hook call if hook is caller.snap b/.forge-snapshots/swap skips hook call if hook is caller.snap
@@ -1 +1 @@
-217673
+217547
diff --git a/.forge-snapshots/swap with dynamic fee.snap b/.forge-snapshots/swap with dynamic fee.snap
@@ -1 +1 @@
-145425
+145362
diff --git a/.forge-snapshots/swap with hooks.snap b/.forge-snapshots/swap with hooks.snap
@@ -1 +1 @@
-138665
+138602
diff --git a/.forge-snapshots/swap with lp fee and protocol fee.snap b/.forge-snapshots/swap with lp fee and protocol fee.snap
@@ -1 +1 @@
-176623
+176594
diff --git a/.forge-snapshots/swap with return dynamic fee.snap b/.forge-snapshots/swap with return dynamic fee.snap
@@ -1 +1 @@
-152437
+152374
diff --git a/.forge-snapshots/update dynamic fee in before swap.snap b/.forge-snapshots/update dynamic fee in before swap.snap
@@ -1 +1 @@
-155039
+155010
diff --git a/src/PoolManager.sol b/src/PoolManager.sol
@@ -156,12 +156,13 @@ contract PoolManager is IPoolManager, ProtocolFees, NoDelegateCall, ERC6909Claim
         bytes calldata hookData
     ) external override onlyWhenUnlocked returns (BalanceDelta callerDelta, BalanceDelta feesAccrued) {
         PoolId id = key.toId();
-        _checkPoolInitialized(id);
+        Pool.State storage pool = _getPool(id);
+        pool.checkPoolInitialized();
 
         key.hooks.beforeModifyLiquidity(key, params, hookData);
 
         BalanceDelta principalDelta;
-        (principalDelta, feesAccrued) = _pools[id].modifyLiquidity(
+        (principalDelta, feesAccrued) = pool.modifyLiquidity(
             Pool.ModifyLiquidityParams({
                 owner: msg.sender,
                 tickLower: params.tickLower,
@@ -193,9 +194,9 @@ contract PoolManager is IPoolManager, ProtocolFees, NoDelegateCall, ERC6909Claim
         returns (BalanceDelta swapDelta)
     {
         if (params.amountSpecified == 0) SwapAmountCannotBeZero.selector.revertWith();
-
         PoolId id = key.toId();
-        _checkPoolInitialized(id);
+        Pool.State storage pool = _getPool(id);
+        pool.checkPoolInitialized();
 
         BeforeSwapDelta beforeSwapDelta;
         {
@@ -205,6 +206,7 @@ contract PoolManager is IPoolManager, ProtocolFees, NoDelegateCall, ERC6909Claim
 
             // execute swap, account protocol fees, and emit swap event
             swapDelta = _swap(
+                pool,
                 id,
                 Pool.SwapParams({
                     tickSpacing: key.tickSpacing,
@@ -227,9 +229,11 @@ contract PoolManager is IPoolManager, ProtocolFees, NoDelegateCall, ERC6909Claim
     }
 
     /// @notice Internal swap function to execute a swap, take protocol fees on input token, and emit the swap event
-    function _swap(PoolId id, Pool.SwapParams memory params, Currency inputCurrency) internal returns (BalanceDelta) {
-        (BalanceDelta delta, uint256 feeForProtocol, uint24 swapFee, Pool.SwapState memory state) =
-            _pools[id].swap(params);
+    function _swap(Pool.State storage pool, PoolId id, Pool.SwapParams memory params, Currency inputCurrency)
+        internal
+        returns (BalanceDelta)
+    {
+        (BalanceDelta delta, uint256 feeForProtocol, uint24 swapFee, Pool.SwapState memory state) = pool.swap(params);
 
         // The fee is on the input currency.
         if (feeForProtocol > 0) _updateProtocolFees(inputCurrency, feeForProtocol);
@@ -248,12 +252,12 @@ contract PoolManager is IPoolManager, ProtocolFees, NoDelegateCall, ERC6909Claim
         onlyWhenUnlocked
         returns (BalanceDelta delta)
     {
-        PoolId id = key.toId();
-        _checkPoolInitialized(id);
+        Pool.State storage pool = _getPool(key.toId());
+        pool.checkPoolInitialized();
 
         key.hooks.beforeDonate(key, amount0, amount1, hookData);
 
-        delta = _pools[id].donate(amount0, amount1);
+        delta = pool.donate(amount0, amount1);
 
         _accountPoolBalanceDelta(key, delta, msg.sender);
 
@@ -330,11 +334,6 @@ contract PoolManager is IPoolManager, ProtocolFees, NoDelegateCall, ERC6909Claim
         _accountDelta(key.currency1, delta.amount1(), target);
     }
 
-    /// @notice Checks if a given pool has been initialized
-    function _checkPoolInitialized(PoolId id) internal view {
-        if (_pools[id].isNotInitialized()) PoolNotInitialized.selector.revertWith();
-    }
-
     /// @notice implementation of the _getPool function defined in ProtocolFees
     function _getPool(PoolId id) internal view override returns (Pool.State storage) {
         return _pools[id];

diff --git a/src/libraries/Pool.sol b/src/libraries/Pool.sol
@@ -110,14 +110,13 @@ library Pool {
     }
 
     function setProtocolFee(State storage self, uint24 protocolFee) internal {
-        if (self.isNotInitialized()) PoolNotInitialized.selector.revertWith();
-
+        self.checkPoolInitialized();
         self.slot0 = self.slot0.setProtocolFee(protocolFee);
     }
 
     /// @notice Only dynamic fee pools may update the lp fee.
     function setLPFee(State storage self, uint24 lpFee) internal {
-        if (self.isNotInitialized()) PoolNotInitialized.selector.revertWith();
+        self.checkPoolInitialized();
         self.slot0 = self.slot0.setLpFee(lpFee);
     }
 
@@ -572,8 +571,9 @@ library Pool {
         }
     }
 
-    function isNotInitialized(State storage self) internal view returns (bool) {
-        return self.slot0.sqrtPriceX96() == 0;
+    /// @notice Reverts if the given pool has not been initialized
+    function checkPoolInitialized(State storage self) internal view {
+        if (self.slot0.sqrtPriceX96() == 0) PoolNotInitialized.selector.revertWith();
     }
 
     /// @notice Clears tick data