zama-ai · melanciani · Jan 5, 2026 · Dec 30, 2025
@@ -276,7 +276,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;
     }
@@ -285,47 +285,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;
     }
 
@@ -370,7 +370,7 @@ contract OperatorRewarder {
      * @param to The address to transfer the tokens to.
      * @param amount The amount of tokens to transfer.
      */
-    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));
@@ -454,7 +454,7 @@ contract OperatorRewarder {
      * - the total rewards paid to the delegators (total paid rewards)
      * @return Total assets plus earned rewards plus paid rewards.
      */
-    function _totalAssetsPlusPaidRewards() internal view returns (uint256) {
+    function _totalAssetsPlusPaidRewards() internal view virtual returns (uint256) {
         return
             token().balanceOf(address(this)) +
             (isShutdown() ? 0 : protocolStaking().earned(address(operatorStaking()))) +
@@ -474,7 +474,7 @@ contract OperatorRewarder {
      * @param totalAssetsPlusPaidRewards The total assets plus earned rewards plus paid rewards.
      * @return Amount of unpaid fee.
      */
-    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;
     }
@@ -492,7 +492,7 @@ contract OperatorRewarder {
      * - paid virtual rewards: a pool of "virtual" rewards that account for changes in the weight distribution
      *  Note: the `mulDiv` rounds down: floor(totalRewards * share / total)
      */
-    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);
     }

@@ -119,7 +119,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();
@@ -414,7 +414,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));
@@ -445,7 +445,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.
@@ -471,7 +471,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
         }

@@ -135,7 +135,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);
@@ -150,7 +150,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);
 
@@ -168,7 +168,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();
@@ -201,7 +201,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);
@@ -215,7 +215,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_);
     }
 
@@ -225,7 +225,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);
     }
 
@@ -235,7 +235,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);
     }
 
@@ -244,7 +244,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_);
     }
 
@@ -254,7 +254,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);
@@ -265,7 +265,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
@@ -278,7 +278,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
      * @dev Returns the staking token which is used for staking and rewards.
      * @return The address of the staking token.
      */
-    function stakingToken() public view returns (address) {
+    function stakingToken() public view virtual returns (address) {
         return _getProtocolStakingStorage()._stakingToken;
     }
 
@@ -287,23 +287,23 @@ 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.
      * @return The total staked weight of all eligible accounts.
      */
-    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.
      * @return The unstake cooldown period in seconds.
      */
-    function unstakeCooldownPeriod() public view returns (uint256) {
+    function unstakeCooldownPeriod() public view virtual returns (uint256) {
         return _getProtocolStakingStorage()._unstakeCooldownPeriod;
     }
 
@@ -322,7 +322,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;
     }
 
@@ -331,7 +331,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;
     }
@@ -341,11 +341,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));
@@ -354,15 +354,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_;
 
@@ -373,7 +373,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();
@@ -382,7 +382,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;
@@ -404,7 +404,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)) {
@@ -420,21 +420,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
         }

@@ -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
-        }
-    }
 }