abstract with ShareManager

src/contracts/core/DelegationManager.sol

@@ -560,9 +560,16 @@ contract DelegationManager is
         require(tokens.length == withdrawal.strategies.length, InputArrayLengthMismatch());
         require(msg.sender == withdrawal.withdrawer, WithdrawerNotCaller());
         bytes32 withdrawalRoot = calculateWithdrawalRoot(withdrawal);
-        bool isLegacyWithdrawal = false;
-
         require(pendingWithdrawals[withdrawalRoot], WithdrawalNotQueued());
+
+        // read delegated operator's totalMagnitudes at time of withdrawal to scale shares again if any slashing has occurred
+        // during withdrawal delay period
+        uint64[] memory totalMagnitudes = allocationManager.getTotalMagnitudesAtTimestamp({
+            operator: withdrawal.delegatedTo,
+            strategies: withdrawal.strategies,
+            timestamp: withdrawal.startTimestamp + MIN_WITHDRAWAL_DELAY
+        });
+
         if (withdrawal.startTimestamp < LEGACY_WITHDRAWALS_TIMESTAMP) {
             // this is a legacy M2 withdrawal using blocknumbers. We use the LEGACY_WITHDRAWALS_TIMESTAMP to check
             // if the withdrawal is a legacy withdrawal or not. It would take up to 600+ years for the blocknumber
@@ -571,64 +578,33 @@ contract DelegationManager is
                 withdrawal.startTimestamp + LEGACY_MIN_WITHDRAWAL_DELAY_BLOCKS <= block.number,
                 WithdrawalDelayNotElapsed()
             );
-            isLegacyWithdrawal = true;
         } else {
             // this is a post Slashing release withdrawal using timestamps
             require(withdrawal.startTimestamp + MIN_WITHDRAWAL_DELAY <= block.timestamp, WithdrawalDelayNotElapsed());
         }
 
-        // read delegated operator's totalMagnitudes at time of withdrawal to scale shares again if any slashing has occurred
-        // during withdrawal delay period
-        uint64[] memory totalMagnitudes = allocationManager.getTotalMagnitudesAtTimestamp({
-            operator: withdrawal.delegatedTo,
-            strategies: withdrawal.strategies,
-            timestamp: withdrawal.startTimestamp + MIN_WITHDRAWAL_DELAY
-        });
-
         for (uint256 i = 0; i < withdrawal.strategies.length; i++) {
-            uint256 sharesToWithdraw;
-            if (isLegacyWithdrawal) {
-                // This is a legacy M2 withdrawal. There is no slashing applied to the withdrawn shares.
-                sharesToWithdraw = withdrawal.scaledShares[i];
-            } else {
-                // Take already scaled staker shares and scale again according to current operator totalMagnitude
-                // This is because the totalMagnitude may have changed since withdrawal was queued and the staker shares
-                // are still susceptible to slashing
-                sharesToWithdraw =
-                    SlashingLib.calculateSharesToCompleteWithdraw(withdrawal.scaledShares[i], totalMagnitudes[i]);
-            }
-
+            IShareManager shareManager = _getShareManager(withdrawal.strategies[i]);
+            uint256 sharesToWithdraw = SlashingLib.calculateSharesToCompleteWithdraw(withdrawal.scaledShares[i], totalMagnitudes[i]);
             if (receiveAsTokens) {
                 // withdraws the underlying token of the strategy to the staker
 
                 // Withdraws `shares` in `strategy` to `withdrawer`. If the shares are virtual beaconChainETH shares,
                 // then a call is ultimately forwarded to the `staker`s EigenPod; otherwise a call is ultimately forwarded
                 // to the `strategy` with info on the `token`.
-                if (withdrawal.strategies[i] == beaconChainETHStrategy) {
-                    eigenPodManager.withdrawSharesAsTokens({
-                        podOwner: withdrawal.staker,
-                        destination: msg.sender,
-                        shares: sharesToWithdraw
-                    });
-                } else {
-                    strategyManager.withdrawSharesAsTokens({
-                        recipient: msg.sender,
-                        strategy: withdrawal.strategies[i],
-                        shares: sharesToWithdraw,
-                        token: tokens[i]
-                    });
-                }
+                shareManager.withdrawSharesAsTokens({
+                    recipient: withdrawal.staker,
+                    strategy: withdrawal.strategies[i],
+                    shares: sharesToWithdraw,
+                    token: tokens[i]
+                });
             } else {
                 // adds the withdrawable shares back to the stakers account
 
                 // Award shares back in StrategyManager/EigenPodManager.
-                if (withdrawal.strategies[i] == beaconChainETHStrategy) {
-                    // Award shares back in EigenPodManager.
-                    // big fukn TODO: refactor EPM to increaseDelegatedShares if applicable. this is completely broken as is
-                    eigenPodManager.addShares({podOwner: withdrawal.staker, shares: sharesToWithdraw});
-                } else {
-                    strategyManager.addShares(msg.sender, tokens[i], withdrawal.strategies[i], sharesToWithdraw);
-                }
+
+                // big fukn TODO: refactor EPM to increaseDelegatedShares if applicable. this is completely broken as is
+                strategyManager.addShares(msg.sender, tokens[i], withdrawal.strategies[i], sharesToWithdraw);
             }
         }
 
@@ -704,16 +680,12 @@ contract DelegationManager is
         // Remove shares from staker and operator
         // Each of these operations fail if we attempt to remove more shares than exist
         for (uint256 i = 0; i < strategies.length; ++i) {
+            IShareManager shareManager = _getShareManager(strategies[i]);
+
             // check sharesToWithdraw is valid
             // TODO maybe have a getter to get totalShares for all strategies, like getDelegatableShares
             // but for inputted strategies
-            uint256 totalShares;
-            if (strategies[i] == beaconChainETHStrategy) {
-                int256 podShares = eigenPodManager.podOwnerShares(staker);
-                totalShares = podShares <= 0 ? 0 : uint256(podShares);
-            } else {
-                totalShares = strategyManager.stakerStrategyShares(staker, strategies[i]);
-            }
+            uint256 totalShares = shareManager.stakerStrategyShares(staker, strategies[i]);
             uint256 stakerScalingFactor = _getStakerScalingFactor(staker, strategies[i]);
             require(
                 sharesToWithdraw[i]
@@ -745,20 +717,11 @@ contract DelegationManager is
                     totalMagnitude: totalMagnitudes[i]
                 });
             }
-
+            
             // Remove active shares from EigenPodManager/StrategyManager
-            if (strategies[i] == beaconChainETHStrategy) {
-                /**
-                 * This call will revert if it would reduce the Staker's virtual beacon chain ETH shares below zero.
-                 * This behavior prevents a Staker from queuing a withdrawal which improperly removes excessive
-                 * shares from the operator to whom the staker is delegated.
-                 * It will also revert if the share amount being withdrawn is not a whole Gwei amount.
-                 */
-                eigenPodManager.removeShares(staker, sharesToDecrement);
-            } else {
-                // this call will revert if `scaledShares[i]` exceeds the Staker's current shares in `strategies[i]`
-                strategyManager.removeShares(staker, strategies[i], sharesToDecrement);
-            }
+            // EigenPodManager: this call will revert if it would reduce the Staker's virtual beacon chain ETH shares below zero
+            // StrategyManager: this call will revert if `scaledShares[i]` exceeds the Staker's current shares in `strategies[i]`
+            shareManager.removeShares(staker, strategies[i], sharesToDecrement);
         }
 
         // Create queue entry and increment withdrawal nonce
@@ -801,6 +764,10 @@ contract DelegationManager is
         return currStakerScalingFactor;
     }
 
+    function _getShareManager(IStrategy strategy) internal view returns (IShareManager) {
+        return strategy == beaconChainETHStrategy ? IShareManager(address(eigenPodManager)) : IShareManager(address(strategyManager));
+    }
+
     /**
      *
      *                         VIEW FUNCTIONS
@@ -885,13 +852,9 @@ contract DelegationManager is
     ) external view returns (uint256[] memory shares) {
         address operator = delegatedTo[staker];
         for (uint256 i = 0; i < strategies.length; ++i) {
+            IShareManager shareManager = _getShareManager(strategies[i]);
             // 1. read strategy shares
-            if (strategies[i] == beaconChainETHStrategy) {
-                int256 podShares = eigenPodManager.podOwnerShares(staker);
-                shares[i] = podShares <= 0 ? 0 : uint256(podShares);
-            } else {
-                shares[i] = strategyManager.stakerStrategyShares(staker, strategies[i]);
-            }
+            shares[i] = shareManager.stakerStrategyShares(staker, strategies[i]);
 
             // 2. if the staker is delegated, actual withdrawable shares can be different from what is stored
             // in the StrategyManager/EigenPodManager because they could have been slashed

src/contracts/core/DelegationManagerStorage.sol

@@ -40,16 +40,18 @@ abstract contract DelegationManagerStorage is IDelegationManager {
 
     /// @notice The AVSDirectory contract for EigenLayer
     IAVSDirectory public immutable avsDirectory;
+
+    // TODO: Switch these to ShareManagers, but this breaks a lot of tests
 
     /// @notice The StrategyManager contract for EigenLayer
     IStrategyManager public immutable strategyManager;
 
-    /// @notice The Slasher contract for EigenLayer
-    ISlasher public immutable slasher;
-
     /// @notice The EigenPodManager contract for EigenLayer
     IEigenPodManager public immutable eigenPodManager;
 
+    /// @notice The Slasher contract for EigenLayer
+    ISlasher public immutable slasher;
+
     /// @notice The AllocationManager contract for EigenLayer
     IAllocationManager public immutable allocationManager;
 

src/contracts/interfaces/IEigenPodManager.sol

@@ -5,6 +5,7 @@ import "@openzeppelin/contracts/proxy/beacon/IBeacon.sol";
 import "./IETHPOSDeposit.sol";
 import "./IStrategyManager.sol";
 import "./IEigenPod.sol";
+import "./IShareManager.sol";
 import "./IPausable.sol";
 import "./ISlasher.sol";
 import "./IStrategy.sol";
@@ -14,7 +15,7 @@ import "./IStrategy.sol";
  * @author Layr Labs, Inc.
  * @notice Terms of Service: https://docs.eigenlayer.xyz/overview/terms-of-service
  */
-interface IEigenPodManager is IPausable {
+interface IEigenPodManager is IShareManager, IPausable {
     /// @dev Thrown when caller is not a EigenPod.
     error OnlyEigenPod();
     /// @dev Thrown when caller is not DelegationManager.
@@ -25,6 +26,8 @@ interface IEigenPodManager is IPausable {
     error SharesNotMultipleOfGwei();
     /// @dev Thrown when shares would result in a negative integer.
     error SharesNegative();
+    /// @dev Thrown when the strategy is not the beaconChainETH strategy.
+    error InvalidStrategy();
 
     /// @notice Emitted to notify the deployment of an EigenPod
     event PodDeployed(address indexed eigenPod, address indexed podOwner);
@@ -118,34 +121,4 @@ interface IEigenPodManager is IPausable {
 
     /// @notice returns canonical, virtual beaconChainETH strategy
     function beaconChainETHStrategy() external view returns (IStrategy);
-
-    /**
-     * @notice Used by the DelegationManager to remove a pod owner's shares while they're in the withdrawal queue.
-     * Simply decreases the `podOwner`'s shares by `shares`, down to a minimum of zero.
-     * @dev This function reverts if it would result in `podOwnerShares[podOwner]` being less than zero, i.e. it is forbidden for this function to
-     * result in the `podOwner` incurring a "share deficit". This behavior prevents a Staker from queuing a withdrawal which improperly removes excessive
-     * shares from the operator to whom the staker is delegated.
-     * @dev Reverts if `shares` is not a whole Gwei amount
-     */
-    function removeShares(address podOwner, uint256 shares) external;
-
-    /**
-     * @notice Increases the `podOwner`'s shares by `shares`, paying off deficit if possible.
-     * Used by the DelegationManager to award a pod owner shares on exiting the withdrawal queue
-     * @dev Returns the number of shares added to `podOwnerShares[podOwner]` above zero, which will be less than the `shares` input
-     * in the event that the podOwner has an existing shares deficit (i.e. `podOwnerShares[podOwner]` starts below zero).
-     * Also returns existingPodShares prior to adding shares, this is returned as 0 if the existing podOwnerShares is negative
-     * @dev Reverts if `shares` is not a whole Gwei amount
-     */
-    function addShares(
-        address podOwner,
-        uint256 shares
-    ) external returns (uint256 increaseInDelegateableShares, uint256 existingPodShares);
-
-    /**
-     * @notice Used by the DelegationManager to complete a withdrawal, sending tokens to some destination address
-     * @dev Prioritizes decreasing the podOwner's share deficit, if they have one
-     * @dev Reverts if `shares` is not a whole Gwei amount
-     */
-    function withdrawSharesAsTokens(address podOwner, address destination, uint256 shares) external;
-}
+}

src/contracts/interfaces/IShareManager.sol

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.27;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "./IStrategy.sol";
+
+/**
+ * @title Interface for a `IShareManager` contract.
+ * @author Layr Labs, Inc.
+ * @notice Terms of Service: https://docs.eigenlayer.xyz/overview/terms-of-service
+ * @notice This contract is used by the DelegationManager as a unified interface to interact with the EigenPodManager and StrategyManager
+ */
+interface IShareManager {
+    /// @notice Used by the DelegationManager to remove a Staker's shares from a particular strategy when entering the withdrawal queue
+    /// @dev strategy must be beaconChainETH when talking to the EigenPodManager
+    function removeShares(address staker, IStrategy strategy, uint256 shares) external;
+
+    /// @notice Used by the DelegationManager to award a Staker some shares that have passed through the withdrawal queue
+    /// @dev strategy must be beaconChainETH when talking to the EigenPodManager
+    /// @dev token is not validated when talking to the EigenPodManager
+    function addShares(address staker, IERC20 token, IStrategy strategy, uint256 shares) external;
+
+    /// @notice Used by the DelegationManager to convert withdrawn descaled shares to tokens and send them to a recipient
+    /// @dev strategy must be beaconChainETH when talking to the EigenPodManager
+    /// @dev token is not validated when talking to the EigenPodManager
+    function withdrawSharesAsTokens(address recipient, IStrategy strategy, uint256 shares, IERC20 token) external;
+
+    /// @notice Returns the current shares of `user` in `strategy`
+    /// @dev strategy must be beaconChainETH when talking to the EigenPodManager
+    /// @dev returns 0 if the user has negative shares
+    function stakerStrategyShares(address user, IStrategy strategy) external view returns (uint256 shares);
+}

src/contracts/interfaces/IStrategyManager.sol

@@ -2,6 +2,7 @@
 pragma solidity >=0.5.0;
 
 import "./IStrategy.sol";
+import "./IShareManager.sol";
 import "./ISlasher.sol";
 import "./IDelegationManager.sol";
 import "./IEigenPodManager.sol";
@@ -12,7 +13,7 @@ import "./IEigenPodManager.sol";
  * @notice Terms of Service: https://docs.eigenlayer.xyz/overview/terms-of-service
  * @notice See the `StrategyManager` contract itself for implementation details.
  */
-interface IStrategyManager {
+interface IStrategyManager is IShareManager {
     /// @dev Thrown when total strategies deployed exceeds max.
     error MaxStrategiesExceeded();
     /// @dev Thrown when two array parameters have mismatching lengths.
@@ -95,18 +96,6 @@ interface IStrategyManager {
         bytes memory signature
     ) external returns (uint256 shares);
 
-    /// @notice Used by the DelegationManager to remove a Staker's shares from a particular strategy when entering the withdrawal queue
-    function removeShares(address staker, IStrategy strategy, uint256 shares) external;
-
-    /// @notice Used by the DelegationManager to award a Staker some shares that have passed through the withdrawal queue
-    function addShares(address staker, IERC20 token, IStrategy strategy, uint256 shares) external;
-
-    /// @notice Used by the DelegationManager to convert withdrawn descaled shares to tokens and send them to a recipient
-    function withdrawSharesAsTokens(address recipient, IStrategy strategy, uint256 shares, IERC20 token) external;
-
-    /// @notice Returns the current shares of `user` in `strategy`
-    function stakerStrategyShares(address user, IStrategy strategy) external view returns (uint256 shares);
-
     /**
      * @notice Get all details on the staker's deposits and corresponding shares
      * @return (staker's strategies, shares in these strategies)