refactor for external call and code reuse

Uniswap · Jun 27, 2024 · 29e7bde · 29e7bde
1 parent 93f012e
commit 29e7bde
Show file tree

Hide file tree

Showing 11 changed files with 45 additions and 127 deletions.
diff --git a/.forge-snapshots/autocompound_exactUnclaimedFees.snap b/.forge-snapshots/autocompound_exactUnclaimedFees.snap
@@ -1 +1 @@
-261480
+262398
diff --git a/.forge-snapshots/autocompound_exactUnclaimedFees_exactCustodiedFees.snap b/.forge-snapshots/autocompound_exactUnclaimedFees_exactCustodiedFees.snap
@@ -1 +1 @@
-193853
+194771
diff --git a/.forge-snapshots/autocompound_excessFeesCredit.snap b/.forge-snapshots/autocompound_excessFeesCredit.snap
@@ -1 +1 @@
-282019
+282937
diff --git a/.forge-snapshots/decreaseLiquidity_erc20.snap b/.forge-snapshots/decreaseLiquidity_erc20.snap
@@ -1 +1 @@
-191804
+193172
diff --git a/.forge-snapshots/decreaseLiquidity_erc6909.snap b/.forge-snapshots/decreaseLiquidity_erc6909.snap
@@ -1 +1 @@
-170671
+172032
diff --git a/.forge-snapshots/increaseLiquidity_erc20.snap b/.forge-snapshots/increaseLiquidity_erc20.snap
@@ -1 +1 @@
-174244
+175155
diff --git a/.forge-snapshots/increaseLiquidity_erc6909.snap b/.forge-snapshots/increaseLiquidity_erc6909.snap
@@ -1 +1 @@
-149826
+150744
diff --git a/.forge-snapshots/mintWithLiquidity.snap b/.forge-snapshots/mintWithLiquidity.snap
@@ -1 +1 @@
-469640
+600083
diff --git a/contracts/NonfungiblePositionManager.sol b/contracts/NonfungiblePositionManager.sol
@@ -39,24 +39,23 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
         ERC721Permit("Uniswap V4 Positions NFT-V1", "UNI-V4-POS", "1")
     {}
 
-    function unlockAndExecute(bytes[] calldata data) external {
-        // TODO: bubble up the return
-        manager.unlock(abi.encode(LiquidityOperation.EXECUTE, abi.encode(data)));
+    function unlockAndExecute(bytes[] memory data) public returns (BalanceDelta delta) {
+        delta = abi.decode(manager.unlock(abi.encode(data)), (BalanceDelta));
     }
 
-    /// @param data bytes[] - array of abi.encodeWithSelector(<selector>, <arguments>)
-    function _execute(bytes[] memory data) internal override returns (bytes memory) {
+    function _unlockCallback(bytes calldata payload) internal override returns (bytes memory) {
+        bytes[] memory data = abi.decode(payload, (bytes[]));
+
         bool success;
+        bytes memory returnData;
         for (uint256 i; i < data.length; i++) {
             // TODO: bubble up the return
-            (success,) = address(this).call(data[i]);
+            (success, returnData) = address(this).call(data[i]);
             if (!success) revert("EXECUTE_FAILED");
         }
-
         // zeroOut();
 
-        // TODO: return something
-        return new bytes(0);
+        return returnData;
     }
 
     // NOTE: more gas efficient as LiquidityAmounts is used offchain
@@ -77,7 +76,9 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
             _closeCallerDeltas(delta, range.poolKey.currency0, range.poolKey.currency1, recipient, false);
             _closeThisDeltas(thisDelta, range.poolKey.currency0, range.poolKey.currency1);
         } else {
-            delta = _unlockAndIncreaseLiquidity(recipient, range, liquidity, hookData, false);
+            bytes[] memory data = new bytes[](1);
+            data[0] = abi.encodeWithSelector(this.mint.selector, range, liquidity, deadline, recipient, hookData);
+            delta = unlockAndExecute(data);
         }
 
         // mint receipt token
@@ -122,7 +123,9 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
             );
             _closeThisDeltas(thisDelta, tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
         } else {
-            delta = _unlockAndIncreaseLiquidity(tokenPos.owner, tokenPos.range, liquidity, hookData, claims);
+            bytes[] memory data = new bytes[](1);
+            data[0] = abi.encodeWithSelector(this.increaseLiquidity.selector, tokenId, liquidity, hookData, claims);
+            delta = unlockAndExecute(data);
         }
     }
 
@@ -133,8 +136,17 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
     {
         TokenPosition memory tokenPos = tokenPositions[tokenId];
 
-        // TODO: @sauce update once _decreaseLiquidity returns callerDelta/thisDelta
-        delta = _unlockAndDecreaseLiquidity(tokenPos.owner, tokenPos.range, liquidity, hookData, claims);
+        if (manager.isUnlocked()) {
+            delta = _decreaseLiquidity(tokenPos.owner, tokenPos.range, liquidity, hookData);
+            _closeCallerDeltas(
+                delta, tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1, tokenPos.owner, claims
+            );
+            _closeAllDeltas(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
+        } else {
+            bytes[] memory data = new bytes[](1);
+            data[0] = abi.encodeWithSelector(this.decreaseLiquidity.selector, tokenId, liquidity, hookData, claims);
+            delta = unlockAndExecute(data);
+        }
     }
 
     function burn(uint256 tokenId, address recipient, bytes calldata hookData, bool claims)
@@ -163,8 +175,17 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
         returns (BalanceDelta delta)
     {
         TokenPosition memory tokenPos = tokenPositions[tokenId];
-        // TODO: @sauce update once _collect returns callerDelta/thisDel
-        delta = _unlockAndCollect(tokenPos.owner, tokenPos.range, hookData, claims);
+        if (manager.isUnlocked()) {
+            delta = _collect(tokenPos.owner, tokenPos.range, hookData);
+            _closeCallerDeltas(
+                delta, tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1, tokenPos.owner, claims
+            );
+            _closeAllDeltas(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
+        } else {
+            bytes[] memory data = new bytes[](1);
+            data[0] = abi.encodeWithSelector(this.collect.selector, tokenId, recipient, hookData, claims);
+            delta = unlockAndExecute(data);
+        }
     }
 
     function feesOwed(uint256 tokenId) external view returns (uint256 token0Owed, uint256 token1Owed) {

diff --git a/contracts/base/BaseLiquidityManagement.sol b/contracts/base/BaseLiquidityManagement.sol
@@ -64,29 +64,6 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         else if (callerDelta1 > 0) currency1.send(manager, owner, uint128(callerDelta1), claims);
     }
 
-    function _execute(bytes[] memory data) internal virtual returns (bytes memory);
-
-    function _unlockCallback(bytes calldata data) internal override returns (bytes memory) {
-        (LiquidityOperation op, bytes memory args) = abi.decode(data, (LiquidityOperation, bytes));
-        if (op == LiquidityOperation.EXECUTE) {
-            (bytes[] memory payload) = abi.decode(args, (bytes[]));
-            return _execute(payload);
-        } else {
-            (address owner, LiquidityRange memory range, uint256 liquidityChange, bytes memory hookData, bool claims) =
-                abi.decode(args, (address, LiquidityRange, uint256, bytes, bool));
-
-            if (op == LiquidityOperation.INCREASE) {
-                return abi.encode(_increaseLiquidityAndZeroOut(owner, range, liquidityChange, hookData, claims));
-            } else if (op == LiquidityOperation.DECREASE) {
-                return abi.encode(_decreaseLiquidityAndZeroOut(owner, range, liquidityChange, hookData, claims));
-            } else if (op == LiquidityOperation.COLLECT) {
-                return abi.encode(_collectAndZeroOut(owner, range, 0, hookData, claims));
-            } else {
-                return new bytes(0);
-            }
-        }
-    }
-
     function _modifyLiquidity(address owner, LiquidityRange memory range, int256 liquidityChange, bytes memory hookData)
         internal
         returns (BalanceDelta liquidityDelta, BalanceDelta totalFeesAccrued)
@@ -163,20 +140,6 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         position.updateFeeGrowthInside(feeGrowthInside0X128, feeGrowthInside1X128);
     }
 
-    function _increaseLiquidityAndZeroOut(
-        address owner,
-        LiquidityRange memory range,
-        uint256 liquidityToAdd,
-        bytes memory hookData,
-        bool claims
-    ) internal returns (BalanceDelta callerDelta) {
-        BalanceDelta thisDelta;
-        // TODO move callerDelta and thisDelta to transient storage?
-        (callerDelta, thisDelta) = _increaseLiquidity(owner, range, liquidityToAdd, hookData);
-        _closeCallerDeltas(callerDelta, range.poolKey.currency0, range.poolKey.currency1, owner, claims);
-        _closeThisDeltas(thisDelta, range.poolKey.currency0, range.poolKey.currency1);
-    }
-
     // When chaining many actions, this should be called at the very end to close out any open deltas owed to or by this contract for other users on the same range.
     // This is safe because any amounts the caller should not pay or take have already been accounted for in closeCallerDeltas.
     function _closeThisDeltas(BalanceDelta delta, Currency currency0, Currency currency1) internal {
@@ -225,21 +188,6 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         return (tokensOwed, callerDelta, thisDelta);
     }
 
-    function _unlockAndIncreaseLiquidity(
-        address owner,
-        LiquidityRange memory range,
-        uint256 liquidityToAdd,
-        bytes memory hookData,
-        bool claims
-    ) internal returns (BalanceDelta) {
-        return abi.decode(
-            manager.unlock(
-                abi.encode(LiquidityOperation.INCREASE, abi.encode(owner, range, liquidityToAdd, hookData, claims))
-            ),
-            (BalanceDelta)
-        );
-    }
-
     function _decreaseLiquidity(
         address owner,
         LiquidityRange memory range,
@@ -277,33 +225,6 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         return callerFeesAccrued;
     }
 
-    function _decreaseLiquidityAndZeroOut(
-        address owner,
-        LiquidityRange memory range,
-        uint256 liquidityToRemove,
-        bytes memory hookData,
-        bool claims
-    ) internal returns (BalanceDelta delta) {
-        delta = _decreaseLiquidity(owner, range, liquidityToRemove, hookData);
-        _closeCallerDeltas(delta, range.poolKey.currency0, range.poolKey.currency1, owner, claims);
-        _closeAllDeltas(range.poolKey.currency0, range.poolKey.currency1);
-    }
-
-    function _unlockAndDecreaseLiquidity(
-        address owner,
-        LiquidityRange memory range,
-        uint256 liquidityToRemove,
-        bytes memory hookData,
-        bool claims
-    ) internal returns (BalanceDelta) {
-        return abi.decode(
-            manager.unlock(
-                abi.encode(LiquidityOperation.DECREASE, abi.encode(owner, range, liquidityToRemove, hookData, claims))
-            ),
-            (BalanceDelta)
-        );
-    }
-
     function _collect(address owner, LiquidityRange memory range, bytes memory hookData)
         internal
         returns (BalanceDelta)
@@ -332,25 +253,6 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         return callerFeesAccrued;
     }
 
-    function _collectAndZeroOut(address owner, LiquidityRange memory range, uint256, bytes memory hookData, bool claims)
-        internal
-        returns (BalanceDelta delta)
-    {
-        delta = _collect(owner, range, hookData);
-        _closeCallerDeltas(delta, range.poolKey.currency0, range.poolKey.currency1, owner, claims);
-        _closeAllDeltas(range.poolKey.currency0, range.poolKey.currency1);
-    }
-
-    function _unlockAndCollect(address owner, LiquidityRange memory range, bytes memory hookData, bool claims)
-        internal
-        returns (BalanceDelta)
-    {
-        return abi.decode(
-            manager.unlock(abi.encode(LiquidityOperation.COLLECT, abi.encode(owner, range, 0, hookData, claims))),
-            (BalanceDelta)
-        );
-    }
-
     // TODO: I deprecated this bc I liked to see the accounting in line in the top level function... and I like to do all the position updates at once.
     //  can keep but should at at least use the position library in here.
     function _updateFeeGrowth(LiquidityRange memory range, Position storage position)

diff --git a/contracts/interfaces/IBaseLiquidityManagement.sol b/contracts/interfaces/IBaseLiquidityManagement.sol
@@ -21,12 +21,7 @@ interface IBaseLiquidityManagement {
         uint128 tokensOwed1;
     }
 
-    enum LiquidityOperation {
-        INCREASE,
-        DECREASE,
-        COLLECT,
-        EXECUTE
-    }
+    error LockFailure();
 
     /// @notice Fees owed for a given liquidity position. Includes materialized fees and uncollected fees.
     /// @param owner The owner of the liquidity position