From 857da0fb2f90bee74422201fa2ebf17ed63e91c4 Mon Sep 17 00:00:00 2001 From: Michael Sun Date: Fri, 13 Sep 2024 13:43:16 -0400 Subject: [PATCH] refactor: add shares --- src/contracts/core/DelegationManager.sol | 35 +++++++------------ src/contracts/core/StrategyManager.sol | 33 +++++++---------- src/contracts/interfaces/IStrategyManager.sol | 7 +--- src/test/mocks/StrategyManagerMock.sol | 2 +- 4 files changed, 27 insertions(+), 50 deletions(-) diff --git a/src/contracts/core/DelegationManager.sol b/src/contracts/core/DelegationManager.sol index 93319be362..b7a2300bb2 100644 --- a/src/contracts/core/DelegationManager.sol +++ b/src/contracts/core/DelegationManager.sol @@ -636,7 +636,7 @@ contract DelegationManager is ) internal { // Finalize action by converting scaled shares to tokens for each strategy, or // by re-awarding shares in each strategy. - for (uint256 i = 0; i < withdrawal.strategies.length; ++i) { + for (uint256 i = 0; i < withdrawal.strategies.length;) { uint256 sharesToWithdraw; if (isLegacyWithdrawal) { // This is a legacy M2 withdrawal. There is no slashing applied to the withdrawn shares. @@ -666,6 +666,10 @@ contract DelegationManager is token: tokens[i] }); } + + unchecked { + ++i; + } } } @@ -685,7 +689,7 @@ contract DelegationManager is // read delegated operator for scaling and adding shares back if needed address currentOperator = delegatedTo[msg.sender]; - for (uint256 i = 0; i < withdrawal.strategies.length; ++i) { + for (uint256 i = 0; i < withdrawal.strategies.length;) { // store existing shares to calculate new staker scaling factor later uint256 existingShares; uint256 shares; @@ -704,6 +708,8 @@ contract DelegationManager is * Other strategy shares can + will be awarded to the withdrawer. */ if (withdrawal.strategies[i] == beaconChainETHStrategy) { + // TODO: REFACTOR EPM AND NEGATIVE SHARES + address staker = withdrawal.staker; /** * Update shares amount depending upon the returned value. @@ -723,29 +729,12 @@ contract DelegationManager is }); } } else { - existingShares = strategyManager.addShares(msg.sender, tokens[i], withdrawal.strategies[i], shares); - // Similar to `isDelegated` logic - if (currentOperator != address(0)) { - _increaseOperatorScaledShares({ - operator: currentOperator, - // the 'staker' here is the address receiving new shares - staker: msg.sender, - strategy: withdrawal.strategies[i], - shares: shares, - totalMagnitude: totalMagnitudes[i] - }); - } + strategyManager.addShares(msg.sender, tokens[i], withdrawal.strategies[i], shares); } - // update stakers scaling deposit scaling factor - uint256 newStakerScalingFactor = SlashingLib.calculateStakerScalingFactor({ - staker: withdrawal.staker, - currStakerScalingFactor: _getStakerScalingFactor(withdrawal.staker, withdrawal.strategies[i]), - totalMagnitude: totalMagnitudes[i], - existingShares: existingShares, - addedShares: shares - }); - stakerScalingFactors[withdrawal.staker][withdrawal.strategies[i]] = newStakerScalingFactor; + unchecked { + ++i; + } } } diff --git a/src/contracts/core/StrategyManager.sol b/src/contracts/core/StrategyManager.sol index 74517e3a8a..b28d50ef69 100644 --- a/src/contracts/core/StrategyManager.sol +++ b/src/contracts/core/StrategyManager.sol @@ -168,9 +168,8 @@ contract StrategyManager is IERC20 token, IStrategy strategy, uint256 shares - ) external onlyDelegationManager returns (uint256 existingShares) { - uint256 existingShares = _addShares(staker, token, strategy, shares); - return existingShares; + ) external onlyDelegationManager { + _addShares(staker, token, strategy, shares); } /// @notice Used by the DelegationManager to convert withdrawn shares to tokens and send them to a recipient @@ -241,12 +240,7 @@ contract StrategyManager is * delegated shares are tracked, increases the stored share amount in `stakerStrategyShares[staker][strategy]`, and adds `strategy` * to the `staker`'s list of strategies, if it is not in the list already. */ - function _addShares( - address staker, - IERC20 token, - IStrategy strategy, - uint256 shares - ) internal returns (uint256 existingShares) { + function _addShares(address staker, IERC20 token, IStrategy strategy, uint256 shares) internal { // sanity checks on inputs require(staker != address(0), StakerAddressZero()); require(shares != 0, SharesAmountZero()); @@ -258,11 +252,18 @@ contract StrategyManager is } // add the returned shares to their existing shares for this strategy - existingShares = stakerStrategyShares[staker][strategy]; + uint256 existingShares = stakerStrategyShares[staker][strategy]; stakerStrategyShares[staker][strategy] += shares; + // Increase shares delegated to operator, if needed + delegation.increaseDelegatedShares({ + staker: staker, + strategy: strategy, + existingShares: existingShares, + addedShares: shares + }); + emit Deposit(staker, token, strategy, shares); - return existingShares; } /** @@ -287,15 +288,7 @@ contract StrategyManager is uint256 shares = strategy.deposit(token, amount); // add the returned shares to the staker's existing shares for this strategy - uint256 existingShares = _addShares(staker, token, strategy, shares); - - // Increase shares delegated to operator, if needed - delegation.increaseDelegatedShares({ - staker: staker, - strategy: strategy, - existingShares: existingShares, - addedShares: shares - }); + _addShares(staker, token, strategy, shares); return shares; } diff --git a/src/contracts/interfaces/IStrategyManager.sol b/src/contracts/interfaces/IStrategyManager.sol index c2037337ef..5e95e835fd 100644 --- a/src/contracts/interfaces/IStrategyManager.sol +++ b/src/contracts/interfaces/IStrategyManager.sol @@ -99,12 +99,7 @@ interface IStrategyManager { 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 returns (uint256 existingShares); + 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; diff --git a/src/test/mocks/StrategyManagerMock.sol b/src/test/mocks/StrategyManagerMock.sol index 2aec094dfe..656725d412 100644 --- a/src/test/mocks/StrategyManagerMock.sol +++ b/src/test/mocks/StrategyManagerMock.sol @@ -110,7 +110,7 @@ contract StrategyManagerMock is function removeShares(address staker, IStrategy strategy, uint256 shares) external {} - function addShares(address staker, IERC20 token, IStrategy strategy, uint256 shares) external returns (uint256 existingShares) {} + function addShares(address staker, IERC20 token, IStrategy strategy, uint256 shares) external {} function withdrawSharesAsTokens(address recipient, IStrategy strategy, uint256 shares, IERC20 token) external {}