lock by address (Uniswap#181)

* Revert "Revert "rough first pass""

This reverts commit 91efa60.

* fix snaps

* clean up impl

* improve reentrancy tests

* add comment

* snaps

* snaps

* cache length

* beef up tests a bit

* make nonzeroDeltaCount global (Uniswap#241)

* make nonzerodeltacount global

* clean up storage

* explicitly pack Owned storage

make _accountDelta more efficient

* clean up visibility/getters

* name change

* address comment

* improve storage efficiency

* update storage order

* snaps

* remove incorrect comment

* nits

* snaps

* address comments

* snaps

* remove parentLockIndex

Author: NoahZinsmeister

.forge-snapshots/donate gas with 1 token.snap

@@ -1 +1 @@
-131226
+93675

.forge-snapshots/donate gas with 2 tokens.snap

@@ -1 +1 @@
-185905
+150995

.forge-snapshots/gas overhead of no-op lock.snap

@@ -1 +1 @@
-61036
+58509

.forge-snapshots/initialize.snap

@@ -1 +1 @@
-37613
+37658

.forge-snapshots/mint with empty hook.snap

@@ -1 +1 @@
-320400
+298837

.forge-snapshots/mint with native token.snap

@@ -1 +1 @@
-294413
+272850

.forge-snapshots/mint.snap

@@ -1 +1 @@
-313077
+291514

.forge-snapshots/poolExtsloadSlot0.snap

@@ -1 +1 @@
-1151
+1151

.forge-snapshots/poolExtsloadTickInfoStruct.snap

@@ -1 +1 @@
-2785
+2785

.forge-snapshots/simple swap.snap

@@ -1 +1 @@
-67705
+47411

.forge-snapshots/swap against liquidity with native token.snap

@@ -1 +1 @@
-161567
+122677

.forge-snapshots/swap against liquidity.snap

@@ -1 +1 @@
-146274
+107384

.forge-snapshots/swap with hooks.snap

@@ -1 +1 @@
-67680
+47386

.forge-snapshots/swap with native.snap

@@ -1 +1 @@
-67705
+47411

contracts/PoolManager.sol

@@ -6,6 +6,7 @@ import {Pool} from "./libraries/Pool.sol";
 import {SafeCast} from "./libraries/SafeCast.sol";
 import {Position} from "./libraries/Position.sol";
 import {Currency, CurrencyLibrary} from "./libraries/CurrencyLibrary.sol";
+import {LockDataLibrary} from "./libraries/LockDataLibrary.sol";
 
 import {NoDelegateCall} from "./NoDelegateCall.sol";
 import {Owned} from "./Owned.sol";
@@ -29,6 +30,7 @@ contract PoolManager is IPoolManager, Owned, NoDelegateCall, ERC1155, IERC1155Re
     using Hooks for IHooks;
     using Position for mapping(bytes32 => Position.Info);
     using CurrencyLibrary for Currency;
+    using LockDataLibrary for IPoolManager.LockData;
     using Fees for uint24;
 
     /// @inheritdoc IPoolManager
@@ -45,18 +47,11 @@ contract PoolManager is IPoolManager, Owned, NoDelegateCall, ERC1155, IERC1155Re
     IProtocolFeeController public protocolFeeController;
 
     /// @inheritdoc IPoolManager
-    address[] public override lockedBy;
+    IPoolManager.LockData public override lockData;
 
-    /// @member nonzeroDeltaCount The number of entries in the currencyDelta mapping that have a non-zero value
-    /// @member currencyDelta The amount owed to the locker (positive) or owed to the pool (negative) of the currency
-    struct LockState {
-        uint256 nonzeroDeltaCount;
-        mapping(Currency currency => int256) currencyDelta;
-    }
-
-    /// @dev Represents the state of the locker at the given index. Each locker must have net 0 currencies owed before
-    /// releasing their lock. Note this is private because the nested mappings cannot be exposed as a public variable.
-    mapping(uint256 index => LockState) private lockStates;
+    /// @dev Represents the currencies due/owed to each locker.
+    /// Must all net to zero when the last lock is released.
+    mapping(address locker => mapping(Currency currency => int256 currencyDelta)) public currencyDelta;
 
     /// @inheritdoc IPoolManager
     mapping(Currency currency => uint256) public override reservesOf;
@@ -126,6 +121,11 @@ contract PoolManager is IPoolManager, Owned, NoDelegateCall, ERC1155, IERC1155Re
         return pools[id].positions.get(owner, tickLower, tickUpper);
     }
 
+    /// @inheritdoc IPoolManager
+    function getLock(uint256 i) external view override returns (address locker) {
+        return LockDataLibrary.getLock(i);
+    }
+
     /// @inheritdoc IPoolManager
     function initialize(PoolKey memory key, uint160 sqrtPriceX96) external override returns (int24 tick) {
         if (key.fee & Fees.STATIC_FEE_MASK >= 1000000) revert FeeTooLarge();
@@ -155,53 +155,37 @@ contract PoolManager is IPoolManager, Owned, NoDelegateCall, ERC1155, IERC1155Re
         emit Initialize(id, key.currency0, key.currency1, key.fee, key.tickSpacing, key.hooks);
     }
 
-    /// @inheritdoc IPoolManager
-    function lockedByLength() external view returns (uint256) {
-        return lockedBy.length;
-    }
-
-    /// @inheritdoc IPoolManager
-    function getNonzeroDeltaCount(uint256 id) external view returns (uint256) {
-        return lockStates[id].nonzeroDeltaCount;
-    }
-
-    /// @inheritdoc IPoolManager
-    function getCurrencyDelta(uint256 id, Currency currency) external view returns (int256) {
-        return lockStates[id].currencyDelta[currency];
-    }
-
     /// @inheritdoc IPoolManager
     function lock(bytes calldata data) external override returns (bytes memory result) {
-        uint256 id = lockedBy.length;
-        lockedBy.push(msg.sender);
+        lockData.push(msg.sender);
 
         // the caller does everything in this callback, including paying what they owe via calls to settle
-        result = ILockCallback(msg.sender).lockAcquired(id, data);
+        result = ILockCallback(msg.sender).lockAcquired(data);
 
-        unchecked {
-            LockState storage lockState = lockStates[id];
-            if (lockState.nonzeroDeltaCount != 0) revert CurrencyNotSettled();
+        if (lockData.length == 1) {
+            if (lockData.nonzeroDeltaCount != 0) revert CurrencyNotSettled();
+            delete lockData;
+        } else {
+            lockData.pop();
         }
-
-        lockedBy.pop();
     }
 
     function _accountDelta(Currency currency, int128 delta) internal {
         if (delta == 0) return;
 
-        LockState storage lockState = lockStates[lockedBy.length - 1];
-        int256 current = lockState.currencyDelta[currency];
-
+        address locker = lockData.getActiveLock();
+        int256 current = currencyDelta[locker][currency];
         int256 next = current + delta;
+
         unchecked {
             if (next == 0) {
-                lockState.nonzeroDeltaCount--;
+                lockData.nonzeroDeltaCount--;
             } else if (current == 0) {
-                lockState.nonzeroDeltaCount++;
+                lockData.nonzeroDeltaCount++;
             }
         }
 
-        lockState.currencyDelta[currency] = next;
+        currencyDelta[locker][currency] = next;
     }
 
     /// @dev Accumulates a balance change to a map of currency to balance changes
@@ -211,7 +195,7 @@ contract PoolManager is IPoolManager, Owned, NoDelegateCall, ERC1155, IERC1155Re
     }
 
     modifier onlyByLocker() {
-        address locker = lockedBy[lockedBy.length - 1];
+        address locker = lockData.getActiveLock();
         if (msg.sender != locker) revert LockedBy(locker);
         _;
     }

contracts/interfaces/IPoolManager.sol

@@ -142,24 +142,27 @@ interface IPoolManager is IERC1155 {
     /// @notice Returns the reserves for a given ERC20 currency
     function reservesOf(Currency currency) external view returns (uint256);
 
-    /// @notice Initialize the state for a given pool ID
-    function initialize(PoolKey memory key, uint160 sqrtPriceX96) external returns (int24 tick);
+    /// @notice Contains data about pool lockers.
+    struct LockData {
+        /// @notice The current number of active lockers
+        uint128 length;
+        /// @notice The total number of nonzero deltas over all active + completed lockers
+        uint128 nonzeroDeltaCount;
+    }
 
-    /// @notice Represents the stack of addresses that have locked the pool. Each call to #lock pushes the address onto the stack
-    /// @param index The index of the locker, also known as the id of the locker
-    function lockedBy(uint256 index) external view returns (address);
+    /// @notice Returns the locker in the ith position of the locker queue.
+    function getLock(uint256 i) external view returns (address locker);
 
-    /// @notice Getter for the length of the lockedBy array
-    function lockedByLength() external view returns (uint256);
+    /// @notice Returns lock data
+    function lockData() external view returns (uint128 length, uint128 nonzeroDeltaCount);
 
-    /// @notice Returns the count of nonzero deltas for the given locker ID
-    /// @param id The ID of the locker
-    function getNonzeroDeltaCount(uint256 id) external view returns (uint256);
+    /// @notice Initialize the state for a given pool ID
+    function initialize(PoolKey memory key, uint160 sqrtPriceX96) external returns (int24 tick);
 
-    /// @notice Get the current delta for a given currency, and its position in the currencies touched array
-    /// @param id The ID of the locker
+    /// @notice Get the current delta for a locker in the given currency
+    /// @param locker The address of the locker
     /// @param currency The currency for which to lookup the delta
-    function getCurrencyDelta(uint256 id, Currency currency) external view returns (int256);
+    function currencyDelta(address locker, Currency currency) external view returns (int256);
 
     /// @notice All operations go through this function
     /// @param data Any data to pass to the callback, via `ILockCallback(msg.sender).lockCallback(data)`

contracts/interfaces/callback/ILockCallback.sol

@@ -3,8 +3,7 @@ pragma solidity ^0.8.20;
 
 interface ILockCallback {
     /// @notice Called by the pool manager on `msg.sender` when a lock is acquired
-    /// @param id The id of the lock that was acquired
     /// @param data The data that was passed to the call to lock
     /// @return Any data that you want to be returned from the lock call
-    function lockAcquired(uint256 id, bytes calldata data) external returns (bytes memory);
+    function lockAcquired(bytes calldata data) external returns (bytes memory);
 }

contracts/libraries/LockDataLibrary.sol

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.20;
+
+import {IPoolManager} from "../interfaces/IPoolManager.sol";
+
+/// @dev This library manages a custom storage implementation for a queue
+///      that tracks current lockers. The "sentinel" storage slot for this data structure,
+///      always passed in as IPoolManager.LockData storage self, stores not just the current
+///      length of the queue but also the global count of non-zero deltas across all lockers.
+///      The values of the data structure start at OFFSET, and each value is a locker address.
+library LockDataLibrary {
+    uint256 private constant OFFSET = uint256(keccak256("LockData"));
+
+    /// @dev Pushes a locker onto the end of the queue, and updates the sentinel storage slot.
+    function push(IPoolManager.LockData storage self, address locker) internal {
+        // read current value from the sentinel storage slot
+        uint128 length = self.length;
+        unchecked {
+            uint256 indexToWrite = OFFSET + length; // not in assembly because OFFSET is in the library scope
+            /// @solidity memory-safe-assembly
+            assembly {
+                // in the next storage slot, write the locker
+                sstore(indexToWrite, locker)
+            }
+            // update the sentinel storage slot
+            self.length = length + 1;
+        }
+    }
+
+    /// @dev Pops a locker off the end of the queue. Note that no storage gets cleared.
+    function pop(IPoolManager.LockData storage self) internal {
+        unchecked {
+            self.length--;
+        }
+    }
+
+    function getLock(uint256 i) internal view returns (address locker) {
+        unchecked {
+            uint256 position = OFFSET + i; // not in assembly because OFFSET is in the library scope
+            /// @solidity memory-safe-assembly
+            assembly {
+                locker := sload(position)
+            }
+        }
+    }
+
+    function getActiveLock(IPoolManager.LockData storage self) internal view returns (address locker) {
+        return getLock(self.length - 1);
+    }
+}

contracts/test/PoolDonateTest.sol

@@ -38,7 +38,7 @@ contract PoolDonateTest is ILockCallback {
         }
     }
 
-    function lockAcquired(uint256, bytes calldata rawData) external returns (bytes memory) {
+    function lockAcquired(bytes calldata rawData) external returns (bytes memory) {
         require(msg.sender == address(manager));
 
         CallbackData memory data = abi.decode(rawData, (CallbackData));

contracts/test/PoolLockTest.sol

@@ -5,7 +5,7 @@ import {IPoolManager} from "../interfaces/IPoolManager.sol";
 import {ILockCallback} from "../interfaces/callback/ILockCallback.sol";
 
 contract PoolLockTest is ILockCallback {
-    event LockAcquired(uint256 id);
+    event LockAcquired();
 
     IPoolManager manager;
 
@@ -18,8 +18,8 @@ contract PoolLockTest is ILockCallback {
     }
 
     /// @notice Called by the pool manager on `msg.sender` when a lock is acquired
-    function lockAcquired(uint256 id, bytes calldata) external override returns (bytes memory) {
-        emit LockAcquired(id);
+    function lockAcquired(bytes calldata) external override returns (bytes memory) {
+        emit LockAcquired();
         return "";
     }
 }

contracts/test/PoolManagerReentrancyTest.sol

@@ -36,7 +36,7 @@ contract PoolModifyPositionTest is ILockCallback {
         }
     }
 
-    function lockAcquired(uint256, bytes calldata rawData) external returns (bytes memory) {
+    function lockAcquired(bytes calldata rawData) external returns (bytes memory) {
         require(msg.sender == address(manager));
 
         CallbackData memory data = abi.decode(rawData, (CallbackData));