refactor(protocol-contracts): add virtual to staking functions and change private to internal

Author: melanciani

protocol-contracts/staking/contracts/OperatorRewarder.sol

@@ -260,7 +260,7 @@ contract OperatorRewarder {
      * @param receiver The receiver address.
      * @return The claimer address.
      */
-    function claimer(address receiver) public view returns (address) {
+    function claimer(address receiver) public view virtual returns (address) {
         address authorizedClaimer = _authorizedClaimers[receiver];
         return authorizedClaimer == address(0) ? receiver : authorizedClaimer;
     }
@@ -269,47 +269,47 @@ contract OperatorRewarder {
      * @notice Returns the staking token address.
      * @return The IERC20 staking token.
      */
-    function token() public view returns (IERC20) {
+    function token() public view virtual returns (IERC20) {
         return _token;
     }
 
     /**
      * @notice Returns the ProtocolStaking contract address.
      * @return The ProtocolStaking contract.
      */
-    function protocolStaking() public view returns (ProtocolStaking) {
+    function protocolStaking() public view virtual returns (ProtocolStaking) {
         return _protocolStaking;
     }
 
     /**
      * @notice Returns the OperatorStaking contract address.
      * @return The OperatorStaking contract.
      */
-    function operatorStaking() public view returns (OperatorStaking) {
+    function operatorStaking() public view virtual returns (OperatorStaking) {
         return _operatorStaking;
     }
 
     /**
      * @notice Returns true if contract is shutdown.
      * @return True if shutdown, false otherwise.
      */
-    function isShutdown() public view returns (bool) {
+    function isShutdown() public view virtual returns (bool) {
         return _shutdown;
     }
 
     /**
      * @notice Returns the maximum fee in basis points that the beneficiary can set.
      * @return Fee in basis points.
      */
-    function maxFeeBasisPoints() public view returns (uint16) {
+    function maxFeeBasisPoints() public view virtual returns (uint16) {
         return _maxFeeBasisPoints;
     }
 
     /**
      * @notice Returns the fee in basis points.
      * @return Fee in basis points.
      */
-    function feeBasisPoints() public view returns (uint16) {
+    function feeBasisPoints() public view virtual returns (uint16) {
         return _feeBasisPoints;
     }
 
@@ -339,7 +339,7 @@ contract OperatorRewarder {
         return _unpaidFee(_totalAssetsPlusPaidRewards());
     }
 
-    function _doTransferOut(address to, uint256 amount) internal {
+    function _doTransferOut(address to, uint256 amount) internal virtual {
         IERC20 token_ = token();
         if (amount > token_.balanceOf(address(this))) {
             protocolStaking().claimRewards(address(_operatorStaking));
@@ -410,20 +410,20 @@ contract OperatorRewarder {
         _feeBasisPoints = basisPoints;
     }
 
-    function _totalAssetsPlusPaidRewards() internal view returns (uint256) {
+    function _totalAssetsPlusPaidRewards() internal view virtual returns (uint256) {
         return
             token().balanceOf(address(this)) +
             (isShutdown() ? 0 : protocolStaking().earned(address(operatorStaking()))) +
             _totalRewardsPaid;
     }
 
-    function _unpaidFee(uint256 totalAssetsPlusPaidRewards) internal view returns (uint256) {
+    function _unpaidFee(uint256 totalAssetsPlusPaidRewards) internal view virtual returns (uint256) {
         uint256 totalAssetsPlusPaidRewardsDelta = totalAssetsPlusPaidRewards - _lastClaimTotalAssetsPlusPaidRewards;
         return (totalAssetsPlusPaidRewardsDelta * feeBasisPoints()) / 10_000;
     }
 
     /// @dev Compute total allocation based on number of shares and total shares. Must take paid rewards into account after.
-    function _allocation(uint256 share, uint256 total) private view returns (uint256) {
+    function _allocation(uint256 share, uint256 total) internal view virtual returns (uint256) {
         return
             SafeCast.toUint256(SafeCast.toInt256(historicalReward()) + _totalVirtualRewardsPaid).mulDiv(share, total);
     }

protocol-contracts/staking/contracts/OperatorStaking.sol

@@ -101,7 +101,7 @@ contract OperatorStaking is ERC1363Upgradeable, ReentrancyGuardTransient, UUPSUp
         address beneficiary_,
         uint16 initialMaxFeeBasisPoints_,
         uint16 initialFeeBasisPoints_
-    ) public initializer {
+    ) public virtual initializer {
         __ERC20_init(name, symbol);
 
         OperatorStakingStorage storage $ = _getOperatorStakingStorage();
@@ -391,7 +391,7 @@ contract OperatorStaking is ERC1363Upgradeable, ReentrancyGuardTransient, UUPSUp
         return _getOperatorStakingStorage()._operator[controller][operator];
     }
 
-    function _doTransferOut(address to, uint256 amount) internal {
+    function _doTransferOut(address to, uint256 amount) internal virtual {
         IERC20 asset_ = IERC20(asset());
         if (amount > asset_.balanceOf(address(this))) {
             protocolStaking().release(address(this));
@@ -422,7 +422,7 @@ contract OperatorStaking is ERC1363Upgradeable, ReentrancyGuardTransient, UUPSUp
         emit IERC4626.Deposit(caller, receiver, assets, shares);
     }
 
-    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}
+    function _authorizeUpgrade(address newImplementation) internal virtual override onlyOwner {}
 
     function _convertToShares(uint256 assets, Math.Rounding rounding) internal view virtual returns (uint256) {
         // Shares in redemption have not yet received assets, so we need to account for them in the conversion.
@@ -448,7 +448,7 @@ contract OperatorStaking is ERC1363Upgradeable, ReentrancyGuardTransient, UUPSUp
         return 0;
     }
 
-    function _getOperatorStakingStorage() private pure returns (OperatorStakingStorage storage $) {
+    function _getOperatorStakingStorage() internal pure returns (OperatorStakingStorage storage $) {
         assembly {
             $.slot := OPERATOR_STAKING_STORAGE_LOCATION
         }

protocol-contracts/staking/contracts/ProtocolStaking.sol

@@ -91,7 +91,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         address manager,
         uint48 initialUnstakeCooldownPeriod,
         uint256 initialRewardRate
-    ) public initializer {
+    ) public virtual initializer {
         __AccessControlDefaultAdminRules_init(0, governor);
         _grantRole(MANAGER_ROLE, manager);
         _setRoleAdmin(ELIGIBLE_ACCOUNT_ROLE, MANAGER_ROLE);
@@ -106,7 +106,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @dev Stake `amount` tokens from `msg.sender`.
      * @param amount The amount of tokens to stake.
      */
-    function stake(uint256 amount) public {
+    function stake(uint256 amount) public virtual {
         _mint(msg.sender, amount);
         IERC20(stakingToken()).safeTransferFrom(msg.sender, address(this), amount);
 
@@ -124,7 +124,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @param amount The amount of tokens to unstake.
      * @return releaseTime The timestamp when the unstaked tokens can be released.
      */
-    function unstake(uint256 amount) public returns (uint48) {
+    function unstake(uint256 amount) public virtual returns (uint48) {
         _burn(msg.sender, amount);
 
         ProtocolStakingStorage storage $ = _getProtocolStakingStorage();
@@ -157,7 +157,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @dev Claim staking rewards for `account`. Can be called by anyone.
      * @param account The account to claim rewards for.
      */
-    function claimRewards(address account) public {
+    function claimRewards(address account) public virtual {
         uint256 rewards = earned(account);
         if (rewards > 0) {
             _getProtocolStakingStorage()._paid[account] += SafeCast.toInt256(rewards);
@@ -171,7 +171,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @dev Sets the reward rate in tokens per second. Only callable by `MANAGER_ROLE` role.
      * @param rewardRate_ The new reward rate in tokens per second.
      */
-    function setRewardRate(uint256 rewardRate_) public onlyRole(MANAGER_ROLE) {
+    function setRewardRate(uint256 rewardRate_) public virtual onlyRole(MANAGER_ROLE) {
         _setRewardRate(rewardRate_);
     }
 
@@ -181,7 +181,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * `ELIGIBLE_ACCOUNT_ROLE`. By default this is `MANAGER_ROLE`.
      * @param account The account to grant the `ELIGIBLE_ACCOUNT_ROLE` role to.
      */
-    function addEligibleAccount(address account) public {
+    function addEligibleAccount(address account) public virtual {
         grantRole(ELIGIBLE_ACCOUNT_ROLE, account);
     }
 
@@ -191,7 +191,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * `ELIGIBLE_ACCOUNT_ROLE`. By default this is `MANAGER_ROLE`.
      * @param account The account to revoke the `ELIGIBLE_ACCOUNT_ROLE` role from.
      */
-    function removeEligibleAccount(address account) public {
+    function removeEligibleAccount(address account) public virtual {
         revokeRole(ELIGIBLE_ACCOUNT_ROLE, account);
     }
 
@@ -200,7 +200,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * by `MANAGER_ROLE` role. See {unstake} for important notes regarding the cooldown period.
      * @param unstakeCooldownPeriod_ The new unstake cooldown period.
      */
-    function setUnstakeCooldownPeriod(uint48 unstakeCooldownPeriod_) public onlyRole(MANAGER_ROLE) {
+    function setUnstakeCooldownPeriod(uint48 unstakeCooldownPeriod_) public virtual onlyRole(MANAGER_ROLE) {
         _setUnstakeCooldownPeriod(unstakeCooldownPeriod_);
     }
 
@@ -210,7 +210,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @param recipient The recipient that will receive rewards on behalf of `msg.sender` for all future  {claimRewards} calls.
      * A value of `address(0)` indicates that rewards should be sent to `msg.sender`.
      */
-    function setRewardsRecipient(address recipient) public {
+    function setRewardsRecipient(address recipient) public virtual {
         _getProtocolStakingStorage()._rewardsRecipient[msg.sender] = recipient;
 
         emit RewardsRecipientSet(msg.sender, recipient);
@@ -221,7 +221,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @param account The account to check rewards for.
      * @return The earned amount.
      */
-    function earned(address account) public view returns (uint256) {
+    function earned(address account) public view virtual returns (uint256) {
         ProtocolStakingStorage storage $ = _getProtocolStakingStorage();
         uint256 stakedWeight = isEligibleAccount(account) ? weight(balanceOf(account)) : 0;
         // if stakedWeight == 0, there is a risk of totalStakedWeight == 0. To avoid div by 0 just return 0
@@ -231,7 +231,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
     }
 
     /// @dev Returns the staking token which is used for staking and rewards.
-    function stakingToken() public view returns (address) {
+    function stakingToken() public view virtual returns (address) {
         return _getProtocolStakingStorage()._stakingToken;
     }
 
@@ -240,17 +240,17 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @param amount The amount being weighted.
      * @return The staking weight.
      */
-    function weight(uint256 amount) public pure returns (uint256) {
+    function weight(uint256 amount) public pure virtual returns (uint256) {
         return Math.sqrt(amount);
     }
 
     /// @dev Returns the current total staked weight.
-    function totalStakedWeight() public view returns (uint256) {
+    function totalStakedWeight() public view virtual returns (uint256) {
         return _getProtocolStakingStorage()._totalEligibleStakedWeight;
     }
 
     /// @dev Returns the current unstake cooldown period in seconds.
-    function unstakeCooldownPeriod() public view returns (uint256) {
+    function unstakeCooldownPeriod() public view virtual returns (uint256) {
         return _getProtocolStakingStorage()._unstakeCooldownPeriod;
     }
 
@@ -269,7 +269,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @dev Gets the current protocol reward rate in tokens distributed per second.
      * @return The reward rate.
      */
-    function rewardRate() public view returns (uint256) {
+    function rewardRate() public view virtual returns (uint256) {
         return _getProtocolStakingStorage()._rewardRate;
     }
 
@@ -278,7 +278,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @param account The account that earned rewards.
      * @return The rewards recipient.
      */
-    function rewardsRecipient(address account) public view returns (address) {
+    function rewardsRecipient(address account) public view virtual returns (address) {
         address storedRewardsRecipient = _getProtocolStakingStorage()._rewardsRecipient[account];
         return storedRewardsRecipient == address(0) ? account : storedRewardsRecipient;
     }
@@ -288,11 +288,11 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @param account The account being checked for eligibility.
      * @return True if eligible.
      */
-    function isEligibleAccount(address account) public view returns (bool) {
+    function isEligibleAccount(address account) public view virtual returns (bool) {
         return hasRole(ELIGIBLE_ACCOUNT_ROLE, account);
     }
 
-    function _grantRole(bytes32 role, address account) internal override returns (bool) {
+    function _grantRole(bytes32 role, address account) internal virtual override returns (bool) {
         bool success = super._grantRole(role, account);
         if (role == ELIGIBLE_ACCOUNT_ROLE && success) {
             require(account != address(0), InvalidEligibleAccount(account));
@@ -301,15 +301,15 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         return success;
     }
 
-    function _revokeRole(bytes32 role, address account) internal override returns (bool) {
+    function _revokeRole(bytes32 role, address account) internal virtual override returns (bool) {
         bool success = super._revokeRole(role, account);
         if (role == ELIGIBLE_ACCOUNT_ROLE && success) {
             _updateRewards(account, weight(balanceOf(account)), 0);
         }
         return success;
     }
 
-    function _setUnstakeCooldownPeriod(uint48 unstakeCooldownPeriod_) internal {
+    function _setUnstakeCooldownPeriod(uint48 unstakeCooldownPeriod_) internal virtual {
         require(unstakeCooldownPeriod_ != 0 && unstakeCooldownPeriod_ <= 365 days, InvalidUnstakeCooldownPeriod());
         _getProtocolStakingStorage()._unstakeCooldownPeriod = unstakeCooldownPeriod_;
 
@@ -320,7 +320,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @dev Sets the reward rate in tokens per second.
      * @param rewardRate_ The new reward rate in tokens per second.
      */
-    function _setRewardRate(uint256 rewardRate_) internal {
+    function _setRewardRate(uint256 rewardRate_) internal virtual {
         ProtocolStakingStorage storage $ = _getProtocolStakingStorage();
         $._lastUpdateReward = _historicalReward();
         $._lastUpdateTimestamp = Time.timestamp();
@@ -329,7 +329,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         emit RewardRateSet(rewardRate_);
     }
 
-    function _updateRewards(address user, uint256 weightBefore, uint256 weightAfter) internal {
+    function _updateRewards(address user, uint256 weightBefore, uint256 weightAfter) internal virtual {
         ProtocolStakingStorage storage $ = _getProtocolStakingStorage();
         uint256 oldTotalWeight = $._totalEligibleStakedWeight;
         $._totalEligibleStakedWeight = oldTotalWeight - weightBefore + weightAfter;
@@ -351,7 +351,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         }
     }
 
-    function _update(address from, address to, uint256 value) internal override {
+    function _update(address from, address to, uint256 value) internal virtual override {
         // Disable Transfers
         require(from == address(0) || to == address(0), TransferDisabled());
         if (isEligibleAccount(from)) {
@@ -367,21 +367,21 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         super._update(from, to, value);
     }
 
-    function _authorizeUpgrade(address newImplementation) internal override onlyRole(DEFAULT_ADMIN_ROLE) {}
+    function _authorizeUpgrade(address newImplementation) internal virtual override onlyRole(DEFAULT_ADMIN_ROLE) {}
 
-    function _historicalReward() internal view returns (uint256) {
+    function _historicalReward() internal view virtual returns (uint256) {
         ProtocolStakingStorage storage $ = _getProtocolStakingStorage();
         return $._lastUpdateReward + (Time.timestamp() - $._lastUpdateTimestamp) * $._rewardRate;
     }
 
-    function _allocation(uint256 share, uint256 total) private view returns (uint256) {
+    function _allocation(uint256 share, uint256 total) internal view virtual returns (uint256) {
         return
             SafeCast
                 .toUint256(SafeCast.toInt256(_historicalReward()) + _getProtocolStakingStorage()._totalVirtualPaid)
                 .mulDiv(share, total);
     }
 
-    function _getProtocolStakingStorage() private pure returns (ProtocolStakingStorage storage $) {
+    function _getProtocolStakingStorage() internal pure returns (ProtocolStakingStorage storage $) {
         assembly {
             $.slot := PROTOCOL_STAKING_STORAGE_LOCATION
         }

protocol-contracts/staking/contracts/mocks/ProtocolStakingSlashingMock.sol

@@ -12,9 +12,6 @@ contract ProtocolStakingSlashingMock is ProtocolStaking {
     using Checkpoints for Checkpoints.Trace208;
     using SafeERC20 for IERC20;
 
-    bytes32 private constant PROTOCOL_STAKING_STORAGE_LOCATION =
-        0xd955b2342c0487c5e5b5f50f5620ec67dcb16d94462ba5d080d7b7472b67b900;
-
     mapping(address => uint256) private _slashedAmount;
 
     function slash(address account, uint256 amount) public {
@@ -30,21 +27,15 @@ contract ProtocolStakingSlashingMock is ProtocolStaking {
     }
 
     function tokensToReleaseAt(address account, uint48 timestamp) public view virtual returns (uint256) {
-        ProtocolStakingStorage storage $ = __getProtocolStakingStorage();
+        ProtocolStakingStorage storage $ = _getProtocolStakingStorage();
         return $._unstakeRequests[account].upperLookup(timestamp) - $._released[account] - _slashedAmount[account];
     }
 
     function release(address account) public virtual override {
         uint256 amountToRelease = tokensToReleaseAt(account, Time.timestamp());
         if (amountToRelease > 0) {
-            __getProtocolStakingStorage()._released[account] += amountToRelease;
+            _getProtocolStakingStorage()._released[account] += amountToRelease;
             IERC20(stakingToken()).safeTransfer(account, amountToRelease);
         }
     }
-
-    function __getProtocolStakingStorage() private pure returns (ProtocolStaking.ProtocolStakingStorage storage $) {
-        assembly {
-            $.slot := PROTOCOL_STAKING_STORAGE_LOCATION
-        }
-    }
 }