Feat: AllocationManager Storage Simplification (#787)

* feat: cleanup * feat: add helper func * fix: simplification * chore: clean up magnitude info usage and type conversions * refactor: review changes * fix: order struct params by size * fix: correct and simplify calc for slashed pending magnitude * fix: storage gap --------- Co-authored-by: wadealexc <[email protected]> Co-authored-by: clandestine.eth <[email protected]>
Layr-Labs · Oct 1, 2024 · 9d131b1 · 9d131b1
1 parent c9dfe56
commit 9d131b1
Show file tree

Hide file tree

Showing 11 changed files with 209 additions and 290 deletions.
diff --git a/lib/openzeppelin-contracts b/lib/openzeppelin-contracts
diff --git a/lib/openzeppelin-contracts-upgradeable b/lib/openzeppelin-contracts-upgradeable
diff --git a/src/contracts/core/AllocationManager.sol b/src/contracts/core/AllocationManager.sol
diff --git a/src/contracts/core/AllocationManagerStorage.sol b/src/contracts/core/AllocationManagerStorage.sol
@@ -45,17 +45,14 @@ abstract contract AllocationManagerStorage is IAllocationManager {
     /// Note that totalMagnitude is monotonically decreasing and only gets updated upon slashing
     mapping(address => mapping(IStrategy => Snapshots.History)) internal _totalMagnitudeUpdate;
 
-    /// @notice Mapping: operator => strategy => operatorSet (encoded) => snapshotted magnitude
-    mapping(address => mapping(IStrategy => mapping(bytes32 => Snapshots.History))) internal _magnitudeUpdate;
-
     /// @notice Mapping: operator => strategy => OperatorMagnitudeInfo to keep track of info regarding pending magnitude allocations.
-    mapping(address => mapping(IStrategy => OperatorMagnitudeInfo)) public operatorMagnitudeInfo;
+    mapping(address => mapping(IStrategy => FreeMagnitudeInfo)) public operatorFreeMagnitudeInfo;
 
-    /// @notice Mapping: operator => strategy => PendingFreeMagnitude[] to keep track of pending free magnitude from deallocations
-    mapping(address => mapping(IStrategy => PendingFreeMagnitude[])) internal _pendingFreeMagnitude;
+    /// @notice Mapping: operator => strategy => operatorSet (encoded) => MagnitudeInfo
+    mapping(address => mapping(IStrategy => mapping(bytes32 => MagnitudeInfo))) internal _operatorMagnitudeInfo;
 
-    /// @notice Mapping: operator => strategy => operatorSet (encoded) => list of queuedDeallocation indices
-    mapping(address => mapping(IStrategy => mapping(bytes32 => uint256[]))) internal _queuedDeallocationIndices;
+    /// @notice Mapping: operator => strategy => operatorSet[] (encoded) to keep track of pending free magnitude for operatorSet from deallocations
+    mapping(address => mapping(IStrategy => bytes32[])) internal _pendingDeallocationOperatorSets;
 
     /// @notice Mapping: operator => allocation delay (in seconds) for the operator.
     /// This determines how long it takes for allocations to take effect in the future.
@@ -78,5 +75,5 @@ abstract contract AllocationManagerStorage is IAllocationManager {
      * variables without shifting down storage in the inheritance chain.
      * See https://docs.openzeppelin.com/contracts/4.x/upgradeable#storage_gaps
      */
-    uint256[42] private __gap;
+    uint256[44] private __gap;
 }
diff --git a/src/contracts/interfaces/IAllocationManager.sol b/src/contracts/interfaces/IAllocationManager.sol
@@ -25,7 +25,7 @@ interface IAllocationManager is ISignatureUtils {
     /// @dev Thrown when an operator attempts to set their allocation for an operatorSet to the same value
     error SameMagnitude();
     /// @dev Thrown when an allocation is attempted for a given operator when they have pending allocations or deallocations.
-    error PendingAllocationOrDeallocation();
+    error ModificationAlreadyPending();
     /// @dev Thrown when an allocation is attempted that exceeds a given operators total allocatable magnitude.
     error InsufficientAllocatableMagnitude();
     /// @dev Thrown when attempting to use an expired eip-712 signature.
@@ -53,18 +53,30 @@ interface IAllocationManager is ISignatureUtils {
      * @param magnitudeDiff the amount of magnitude to deallocate
      * @param completableTimestamp the timestamp at which the deallocation can be completed, 21 days from when queued
      */
-    struct PendingFreeMagnitude {
-        uint64 magnitudeDiff;
-        uint32 completableTimestamp;
+    // struct PendingFreeMagnitude {
+    //     uint64 magnitudeDiff;
+    //     uint32 completableTimestamp;
+    // }
+
+    /**
+     * @notice struct used for operator magnitude updates. Stored in _operatorMagnitudeInfo mapping
+     * @param currentMagnitude the current magnitude of the operator
+     * @param pendingMagnitudeIdff the pending magnitude difference of the operator
+     * @param effectTimestamp the timestamp at which the pending magnitude will take effect
+     */
+    struct MagnitudeInfo {
+        int128 pendingMagnitudeDiff;
+        uint64 currentMagnitude;
+        uint32 effectTimestamp;
     }
 
     /**
      * @notice Struct containing info regarding free allocatable magnitude.
-     * @param nextPendingFreeMagnitudeIndex The next available update index.
+     * @param nextPendingIndex The next available update index.
      * @param freeMagnitude The total amount of free allocatable magnitude.
      */
-    struct OperatorMagnitudeInfo {
-        uint192 nextPendingFreeMagnitudeIndex;
+    struct FreeMagnitudeInfo {
+        uint192 nextPendingIndex;
         uint64 freeMagnitude;
     }
 
@@ -239,13 +251,14 @@ interface IAllocationManager is ISignatureUtils {
      * @param operator the operator to get the pending deallocations for
      * @param strategy the strategy to get the pending deallocations for
      * @param operatorSets the operatorSets to get the pending deallocations for
-     * @return pendingMagnitudes the latest pending deallocation
+     * @return pendingMagnitudeDiffs the pending deallocation diffs for each operatorSet
+     * @return timestamps the timestamps for each pending dealloction
      */
     function getPendingDeallocations(
         address operator,
         IStrategy strategy,
         OperatorSet[] calldata operatorSets
-    ) external view returns (PendingFreeMagnitude[] memory);
+    ) external view returns (uint64[] memory, uint32[] memory);
 
     /**
      * @notice operator is slashable by operatorSet if currently registered OR last deregistered within 21 days

diff --git a/src/test/events/IRewardsCoordinatorEvents.sol b/src/test/events/IRewardsCoordinatorEvents.sol
@@ -36,14 +36,14 @@ interface IRewardsCoordinatorEvents {
     );
     event ActivationDelaySet(uint32 oldActivationDelay, uint32 newActivationDelay);
     event GlobalCommissionBipsSet(uint16 oldGlobalCommissionBips, uint16 newGlobalCommissionBips);
-    /// @notice emitted when an operator commission is set for a specific OperatorSet
-    event OperatorCommissionUpdated(
-        address indexed operator,
-        OperatorSet indexed operatorSet,
-        IRewardsCoordinator.RewardType rewardType,
-        uint16 newCommissionBips,
-        uint32 effectTimestamp
-    );
+    // /// @notice emitted when an operator commission is set for a specific OperatorSet
+    // event OperatorCommissionUpdated(
+    //     address indexed operator,
+    //     OperatorSet indexed operatorSet,
+    //     IRewardsCoordinator.RewardType rewardType,
+    //     uint16 newCommissionBips,
+    //     uint32 effectTimestamp
+    // );
     event ClaimerForSet(address indexed earner, address indexed oldClaimer, address indexed claimer);
     /// @notice rootIndex is the specific array index of the newly created root in the storage array
     event DistributionRootSubmitted(
@@ -62,13 +62,13 @@ interface IRewardsCoordinatorEvents {
         IERC20 token,
         uint256 claimedAmount
     );
-    /// @notice emitted when an AVS creates a valid OperatorSetRewardsSubmission
-    event OperatorSetRewardCreated(
-        address indexed avs,
-        uint256 indexed submissionNonce,
-        bytes32 indexed rewardsSubmissionHash,
-        IRewardsCoordinator.OperatorSetRewardsSubmission rewardsSubmission
-    );
+    // /// @notice emitted when an AVS creates a valid OperatorSetRewardsSubmission
+    // event OperatorSetRewardCreated(
+    //     address indexed avs,
+    //     uint256 indexed submissionNonce,
+    //     bytes32 indexed rewardsSubmissionHash,
+    //     IRewardsCoordinator.OperatorSetRewardsSubmission rewardsSubmission
+    // );
 
 
 

diff --git a/src/test/mocks/DelegationManagerMock.sol b/src/test/mocks/DelegationManagerMock.sol
@@ -179,7 +179,7 @@ contract DelegationManagerMock is IDelegationManager, Test {
         IStrategy strategy,
         OwnedShares shares
     ) external {
-        strategyManager.addShares(staker, token, strategy, shares);
+        strategyManager.addOwnedShares(staker, strategy, token, shares);
     }
 
     function removeShares(
@@ -198,6 +198,6 @@ contract DelegationManagerMock is IDelegationManager, Test {
         OwnedShares shares,
         IERC20 token
     ) external {
-        strategyManager.withdrawSharesAsTokens(recipient, strategy, shares, token);
+        strategyManager.withdrawSharesAsTokens(recipient, strategy, token, shares);
     }
 }
diff --git a/src/test/mocks/EigenPodManagerMock.sol b/src/test/mocks/EigenPodManagerMock.sol
@@ -19,9 +19,9 @@ contract EigenPodManagerMock is IEigenPodManager, Test, Pausable {
 
     function removeShares(address staker, IStrategy strategy, Shares shares) external {}
 
-    function addShares(address staker, IERC20 token, IStrategy strategy, OwnedShares shares) external {}
+    function addOwnedShares(address staker, IStrategy strategy, IERC20 token, OwnedShares shares) external {}
 
-    function withdrawSharesAsTokens(address recipient, IStrategy strategy, OwnedShares shares, IERC20 token) external {}
+    function withdrawSharesAsTokens(address recipient, IStrategy strategy, IERC20 token, OwnedShares shares) external {}
 
     function stakerStrategyShares(address user, IStrategy strategy) external view returns (Shares shares) {}
 

diff --git a/src/test/mocks/StrategyManagerMock.sol b/src/test/mocks/StrategyManagerMock.sol
@@ -114,9 +114,9 @@ contract StrategyManagerMock is
 
     function removeShares(address staker, IStrategy strategy, Shares shares) external {}
 
-    function addShares(address staker, IERC20 token, IStrategy strategy, OwnedShares shares) external {}
+    function addOwnedShares(address staker, IStrategy strategy, IERC20 token, OwnedShares shares) external {}
 
-    function withdrawSharesAsTokens(address recipient, IStrategy strategy, OwnedShares shares, IERC20 token) external {}
+    function withdrawSharesAsTokens(address recipient, IStrategy strategy, IERC20 token, OwnedShares shares) external {}
 
     /// @notice returns the enshrined beaconChainETH Strategy
     function beaconChainETHStrategy() external view returns (IStrategy) {}

diff --git a/src/test/unit/EigenPodManagerUnit.t.sol b/src/test/unit/EigenPodManagerUnit.t.sol
@@ -194,28 +194,28 @@ contract EigenPodManagerUnitTests_ShareUpdateTests is EigenPodManagerUnitTests {
         cheats.assume(notDelegationManager != address(delegationManagerMock));
         cheats.prank(notDelegationManager);
         cheats.expectRevert("EigenPodManager.onlyDelegationManager: not the DelegationManager");
-        eigenPodManager.addShares(defaultStaker, 0);
+        eigenPodManager.addOwnedShares(defaultStaker, 0);
     }
 
     function test_addShares_revert_podOwnerInputAddressZero() public {
         cheats.prank(address(delegationManagerMock));
-        cheats.expectRevert("EigenPodManager.addShares: podOwner cannot be zero address");
-        eigenPodManager.addShares(address(0), 0);
+        cheats.expectRevert("EigenPodManager.addOwnedShares: podOwner cannot be zero address");
+        eigenPodManager.addOwnedShares(address(0), 0);
     }
 
     function testFuzz_addShares_revert_sharesNegative(int256 shares) public {
         cheats.assume(shares < 0);
         cheats.prank(address(delegationManagerMock));
-        cheats.expectRevert("EigenPodManager.addShares: shares cannot be negative");
-        eigenPodManager.addShares(defaultStaker, uint256(shares));
+        cheats.expectRevert("EigenPodManager.addOwnedShares: shares cannot be negative");
+        eigenPodManager.addOwnedShares(defaultStaker, uint256(shares));
     }
 
     function testFuzz_addShares_revert_sharesNotWholeGwei(uint256 shares) public {
         cheats.assume(int256(shares) >= 0);
         cheats.assume(shares % GWEI_TO_WEI != 0);
         cheats.prank(address(delegationManagerMock));
-        cheats.expectRevert("EigenPodManager.addShares: shares must be a whole Gwei amount");
-        eigenPodManager.addShares(defaultStaker, shares);
+        cheats.expectRevert("EigenPodManager.addOwnedShares: shares must be a whole Gwei amount");
+        eigenPodManager.addOwnedShares(defaultStaker, shares);
     }
 
     function testFuzz_addShares(uint256 shares) public {
@@ -226,7 +226,7 @@ contract EigenPodManagerUnitTests_ShareUpdateTests is EigenPodManagerUnitTests {
 
         // Add shares
         cheats.prank(address(delegationManagerMock));
-        eigenPodManager.addShares(defaultStaker, shares);
+        eigenPodManager.addOwnedShares(defaultStaker, shares);
 
         // Check storage update
         assertEq(eigenPodManager.podOwnerShares(defaultStaker), int256(shares), "Incorrect number of shares added");

diff --git a/src/test/unit/StrategyManagerUnit.t.sol b/src/test/unit/StrategyManagerUnit.t.sol
@@ -979,18 +979,18 @@ contract StrategyManagerUnitTests_addShares is StrategyManagerUnitTests {
     function test_Revert_DelegationManagerModifier() external {
         DelegationManagerMock invalidDelegationManager = new DelegationManagerMock();
         cheats.expectRevert("StrategyManager.onlyDelegationManager: not the DelegationManager");
-        invalidDelegationManager.addShares(strategyManager, address(this), dummyToken, dummyStrat, 1);
+        invalidDelegationManager.addOwnedShares(strategyManager, address(this), dummyToken, dummyStrat, 1);
     }
 
     function testFuzz_Revert_StakerZeroAddress(uint256 amount) external {
         cheats.expectRevert("StrategyManager._addShares: staker cannot be zero address");
-        delegationManagerMock.addShares(strategyManager, address(0), dummyToken, dummyStrat, amount);
+        delegationManagerMock.addOwnedShares(strategyManager, address(0), dummyToken, dummyStrat, amount);
     }
 
     function testFuzz_Revert_ZeroShares(address staker) external filterFuzzedAddressInputs(staker) {
         cheats.assume(staker != address(0));
         cheats.expectRevert("StrategyManager._addShares: shares should not be zero!");
-        delegationManagerMock.addShares(strategyManager, staker, dummyToken, dummyStrat, 0);
+        delegationManagerMock.addOwnedShares(strategyManager, staker, dummyToken, dummyStrat, 0);
     }
 
     function testFuzz_AppendsStakerStrategyList(
@@ -1003,7 +1003,7 @@ contract StrategyManagerUnitTests_addShares is StrategyManagerUnitTests {
         assertEq(sharesBefore, 0, "Staker has already deposited into this strategy");
         assertFalse(_isDepositedStrategy(staker, dummyStrat), "strategy should not be deposited");
 
-        delegationManagerMock.addShares(strategyManager, staker, dummyToken, dummyStrat, amount);
+        delegationManagerMock.addOwnedShares(strategyManager, staker, dummyToken, dummyStrat, amount);
         uint256 stakerStrategyListLengthAfter = strategyManager.stakerStrategyListLength(staker);
         uint256 sharesAfter = strategyManager.stakerStrategyShares(staker, dummyStrat);
         assertEq(
@@ -1028,7 +1028,7 @@ contract StrategyManagerUnitTests_addShares is StrategyManagerUnitTests {
         assertEq(sharesBefore, initialAmount, "Staker has not deposited amount into strategy");
         assertTrue(_isDepositedStrategy(staker, strategy), "strategy should be deposited");
 
-        delegationManagerMock.addShares(strategyManager, staker, dummyToken, dummyStrat, sharesAmount);
+        delegationManagerMock.addOwnedShares(strategyManager, staker, dummyToken, dummyStrat, sharesAmount);
         uint256 stakerStrategyListLengthAfter = strategyManager.stakerStrategyListLength(staker);
         uint256 sharesAfter = strategyManager.stakerStrategyShares(staker, dummyStrat);
         assertEq(
@@ -1077,7 +1077,7 @@ contract StrategyManagerUnitTests_addShares is StrategyManagerUnitTests {
 
         cheats.prank(staker);
         cheats.expectRevert("StrategyManager._addShares: deposit would exceed MAX_STAKER_STRATEGY_LIST_LENGTH");
-        delegationManagerMock.addShares(strategyManager, staker, dummyToken, strategy, amount);
+        delegationManagerMock.addOwnedShares(strategyManager, staker, dummyToken, strategy, amount);
 
         cheats.expectRevert("StrategyManager._addShares: deposit would exceed MAX_STAKER_STRATEGY_LIST_LENGTH");
         strategyManager.depositIntoStrategy(strategy, token, amount);