move decrease and collect to transient storage

Uniswap · Jun 28, 2024 · 5f7e9c1 · 5f7e9c1
1 parent 8ce4ed6
commit 5f7e9c1
Show file tree

Hide file tree

Showing 11 changed files with 48 additions and 18 deletions.
diff --git a/.forge-snapshots/autocompound_exactUnclaimedFees.snap b/.forge-snapshots/autocompound_exactUnclaimedFees.snap
@@ -1 +1 @@
-264058
+265947
diff --git a/.forge-snapshots/autocompound_exactUnclaimedFees_exactCustodiedFees.snap b/.forge-snapshots/autocompound_exactUnclaimedFees_exactCustodiedFees.snap
@@ -1 +1 @@
-196431
+198320
diff --git a/.forge-snapshots/autocompound_excessFeesCredit.snap b/.forge-snapshots/autocompound_excessFeesCredit.snap
@@ -1 +1 @@
-284597
+286486
diff --git a/.forge-snapshots/decreaseLiquidity_erc20.snap b/.forge-snapshots/decreaseLiquidity_erc20.snap
@@ -1 +1 @@
-180891
+184337
diff --git a/.forge-snapshots/decreaseLiquidity_erc6909.snap b/.forge-snapshots/decreaseLiquidity_erc6909.snap
@@ -1 +1 @@
-180903
+184349
diff --git a/.forge-snapshots/increaseLiquidity_erc20.snap b/.forge-snapshots/increaseLiquidity_erc20.snap
@@ -1 +1 @@
-176815
+178704
diff --git a/.forge-snapshots/increaseLiquidity_erc6909.snap b/.forge-snapshots/increaseLiquidity_erc6909.snap
@@ -1 +1 @@
-152404
+154293
diff --git a/.forge-snapshots/mintWithLiquidity.snap b/.forge-snapshots/mintWithLiquidity.snap
@@ -1 +1 @@
-475856
+476807
diff --git a/contracts/NonfungiblePositionManager.sol b/contracts/NonfungiblePositionManager.sol
@@ -119,8 +119,11 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
         TokenPosition memory tokenPos = tokenPositions[tokenId];
 
         if (manager.isUnlocked()) {
-            BalanceDelta thisDelta;
-            (delta, thisDelta) = _increaseLiquidity(tokenPos.owner, tokenPos.range, liquidity, hookData);
+            _increaseLiquidity(tokenPos.owner, tokenPos.range, liquidity, hookData);
+
+            delta = tokenPos.owner.getBalanceDelta(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
+            BalanceDelta thisDelta =
+                address(this).getBalanceDelta(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
 
             // TODO: should be triggered by zeroOut in _execute...
             _closeCallerDeltas(
@@ -143,7 +146,12 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
         TokenPosition memory tokenPos = tokenPositions[tokenId];
 
         if (manager.isUnlocked()) {
-            (delta, thisDelta) = _decreaseLiquidity(tokenPos.owner, tokenPos.range, liquidity, hookData);
+            _decreaseLiquidity(tokenPos.owner, tokenPos.range, liquidity, hookData);
+
+            // TODO: move to zeroOut() in unlockAndExecute()
+            delta = tokenPos.owner.getBalanceDelta(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
+            thisDelta =
+                address(this).getBalanceDelta(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
             _closeCallerDeltas(
                 delta, tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1, tokenPos.owner, claims
             );
@@ -187,8 +195,12 @@ contract NonfungiblePositionManager is INonfungiblePositionManager, BaseLiquidit
     {
         TokenPosition memory tokenPos = tokenPositions[tokenId];
         if (manager.isUnlocked()) {
-            BalanceDelta thisDelta;
-            (delta, thisDelta) = _collect(tokenPos.owner, tokenPos.range, hookData);
+            _collect(tokenPos.owner, tokenPos.range, hookData);
+
+            // TODO: move to zeroOut() in unlockAndExecute()
+            delta = tokenPos.owner.getBalanceDelta(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
+            BalanceDelta thisDelta =
+                address(this).getBalanceDelta(tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1);
             _closeCallerDeltas(
                 delta, tokenPos.range.poolKey.currency0, tokenPos.range.poolKey.currency1, tokenPos.owner, claims
             );

diff --git a/contracts/base/BaseLiquidityManagement.sol b/contracts/base/BaseLiquidityManagement.sol
@@ -43,6 +43,7 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
     using LiquidityDeltaAccounting for BalanceDelta;
     using TransientLiquidityDelta for BalanceDelta;
     using TransientLiquidityDelta for Currency;
+    using TransientLiquidityDelta for address;
 
     mapping(address owner => mapping(LiquidityRangeId rangeId => Position)) public positions;
 
@@ -65,6 +66,8 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
 
         if (callerDelta1 < 0) currency1.settle(manager, owner, uint256(int256(-callerDelta1)), claims);
         else if (callerDelta1 > 0) currency1.take(manager, owner, uint128(callerDelta1), claims);
+
+        owner.close(currency0, currency1);
     }
 
     function _modifyLiquidity(address owner, LiquidityRange memory range, int256 liquidityChange, bytes memory hookData)
@@ -104,8 +107,6 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         // Calculate the portion of the liquidityDelta that is attributable to the caller.
         // We must account for fees that might be owed to other users on the same range.
         (callerDelta, thisDelta) = liquidityDelta.split(callerFeesAccrued, totalFeesAccrued);
-        callerDelta.flush(owner, range.poolKey.currency0, range.poolKey.currency1);
-        thisDelta.flush(address(this), range.poolKey.currency0, range.poolKey.currency1);
 
         // Update position storage, flushing the callerDelta value to tokensOwed first if necessary.
         // If callerDelta > 0, then even after investing callerFeesAccrued, the caller still has some amount to collect that were not added into the position so they are accounted to tokensOwed and removed from the final callerDelta returned.
@@ -119,6 +120,8 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
             (tokensOwed, callerDelta, thisDelta) =
                 _moveCallerDeltaToTokensOwed(false, tokensOwed, callerDelta, thisDelta);
         }
+        callerDelta.flush(owner, range.poolKey.currency0, range.poolKey.currency1);
+        thisDelta.flush(address(this), range.poolKey.currency0, range.poolKey.currency1);
 
         position.addTokensOwed(tokensOwed);
         position.addLiquidity(liquidityToAdd);
@@ -136,6 +139,8 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         // Burn the receipt for tokens owed to this address.
         if (delta0 < 0) currency0.settle(manager, address(this), uint256(int256(-delta0)), true);
         if (delta1 < 0) currency1.settle(manager, address(this), uint256(int256(-delta1)), true);
+
+        address(this).close(currency0, currency1);
     }
 
     function _moveCallerDeltaToTokensOwed(
@@ -157,7 +162,7 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
 
         return (tokensOwed, callerDelta, thisDelta);
     }
-    
+
     /// Any outstanding amounts owed to the caller from the underlying modify call must be collected explicitly with `collect`.
     function _decreaseLiquidity(
         address owner,
@@ -189,6 +194,8 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
             (tokensOwed, callerDelta, thisDelta) =
                 _moveCallerDeltaToTokensOwed(false, tokensOwed, callerDelta, thisDelta);
         }
+        callerDelta.flush(owner, range.poolKey.currency0, range.poolKey.currency1);
+        thisDelta.flush(address(this), range.poolKey.currency0, range.poolKey.currency1);
 
         position.addTokensOwed(tokensOwed);
         position.subtractLiquidity(liquidityToRemove);
@@ -222,6 +229,9 @@ abstract contract BaseLiquidityManagement is IBaseLiquidityManagement, SafeCallb
         callerDelta = callerDelta + tokensOwed;
         thisDelta = thisDelta - tokensOwed;
 
+        callerDelta.flush(owner, range.poolKey.currency0, range.poolKey.currency1);
+        thisDelta.flush(address(this), range.poolKey.currency0, range.poolKey.currency1);
+
         position.clearTokensOwed();
     }
 

diff --git a/contracts/libraries/TransientLiquidityDelta.sol b/contracts/libraries/TransientLiquidityDelta.sol
@@ -7,7 +7,6 @@ import {Currency} from "@uniswap/v4-core/src/types/Currency.sol";
 /// @title a library to store callers' currency deltas in transient storage
 /// @dev this library implements the equivalent of a mapping, as transient storage can only be accessed in assembly
 library TransientLiquidityDelta {
-
     /// @notice calculates which storage slot a delta should be stored in for a given caller and currency
     function _computeSlot(address caller_, Currency currency) internal pure returns (bytes32 hashSlot) {
         assembly {
@@ -17,6 +16,11 @@ library TransientLiquidityDelta {
         }
     }
 
+    function close(address holder, Currency currency0, Currency currency1) internal {
+        setDelta(currency0, holder, 0);
+        setDelta(currency1, holder, 0);
+    }
+
     /// @notice Flush a BalanceDelta into transient storage for a given holder
     function flush(BalanceDelta delta, address holder, Currency currency0, Currency currency1) internal {
         addDelta(currency0, holder, delta.amount0());
@@ -41,7 +45,11 @@ library TransientLiquidityDelta {
         }
     }
 
-    function getBalanceDelta(address holder, Currency currency0, Currency currency1) internal view returns (BalanceDelta delta) {
+    function getBalanceDelta(address holder, Currency currency0, Currency currency1)
+        internal
+        view
+        returns (BalanceDelta delta)
+    {
         delta = toBalanceDelta(getDelta(currency0, holder), getDelta(currency1, holder));
     }