cache non-terminal keys

contracts/PermissionlessNodeRegistry.sol

@@ -21,7 +21,7 @@
 import { IPermissionlessNodeRegistry } from "./interfaces/IPermissionlessNodeRegistry.sol";
 import { IOperatorRewardsCollector } from "./interfaces/IOperatorRewardsCollector.sol";
 
 contract PermissionlessNodeRegistry is
     INodeRegistry,
     IPermissionlessNodeRegistry,
     AccessControlUpgradeable,
@@ -60,6 +60,10 @@
     mapping(uint256 => address) public override nodeELRewardVaultByOperatorId;
     mapping(uint256 => address) public proposedRewardAddressByOperatorId;
     uint256 public maxKeysPerOperator;
+    // mapping of operator Id to cached non-terminal keys count
+    mapping(uint256 => uint64) public operatorNonTerminalKeysCount;
+    // mapping to track if non-terminal keys count has been initialized for an operator
+    mapping(uint256 => bool) public operatorNonTerminalKeysCountInitialized;
 
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
@@ -183,6 +187,10 @@
 
             validatorIdByPubkey[_pubkey[i]] = nextValidatorId;
             validatorIdsByOperatorId[operatorId].push(nextValidatorId);
+            // increment cached count if it was initialized earlier
+            if (operatorNonTerminalKeysCountInitialized[operatorId]) {
+                operatorNonTerminalKeysCount[operatorId]++;
+            }
             emit AddedValidatorKey(msg.sender, _pubkey[i], nextValidatorId);
             nextValidatorId++;
             unchecked {
@@ -275,8 +283,13 @@
             if (!isActiveValidator(validatorId)) {
                 revert UNEXPECTED_STATUS();
             }
+            uint256 operatorId = validatorRegistry[validatorId].operatorId;
             validatorRegistry[validatorId].status = ValidatorStatus.WITHDRAWN;
             validatorRegistry[validatorId].withdrawnBlock = block.number;
+            // decrement cached count if it was initialized earlier
+            if (operatorNonTerminalKeysCountInitialized[operatorId]) {
+                operatorNonTerminalKeysCount[operatorId]--;
+            }
             IValidatorWithdrawalVault(validatorRegistry[validatorId].withdrawVaultAddress).settleFunds();
             emit ValidatorWithdrawn(_pubkeys[i], validatorId);
             unchecked {
@@ -381,6 +394,19 @@
         emit UpdatedStaderConfig(_staderConfig);
     }
 
+    /**
+     * @notice set the cached non-terminal keys count for an operator
+     * @dev only `MANAGER` role can call, used to initialize the cached count
+     * @param _operatorId Id of the operator
+     * @param _nonTerminalKeysCount the non-terminal keys count to set
+     */
+    function setOperatorNonTerminalKeysCount(uint256 _operatorId, uint64 _nonTerminalKeysCount) external {
+        UtilLib.onlyManagerRole(msg.sender, staderConfig);
+        operatorNonTerminalKeysCount[_operatorId] = _nonTerminalKeysCount;
+        operatorNonTerminalKeysCountInitialized[_operatorId] = true;
+        emit OperatorNonTerminalKeysCountSet(_operatorId, _nonTerminalKeysCount);
+    }
+
     /**
      * @notice propose the new reward address of an operator
      * @dev only the existing reward address (msg.sender) can propose
@@ -475,6 +501,14 @@
         uint256 operatorId = operatorIDByAddress[_nodeOperator];
         uint256 validatorCount = getOperatorTotalKeys(operatorId);
         _endIndex = _endIndex > validatorCount ? validatorCount : _endIndex;
+
+        // If cached count is initialized and we're querying the full range (startIndex = 0, endIndex = total),
+        // return the cached value for gas optimization
+        if (operatorNonTerminalKeysCountInitialized[operatorId] && _startIndex == 0 && _endIndex == validatorCount) {
+            return operatorNonTerminalKeysCount[operatorId];
+        }
+
+        // Otherwise, fall back to the loop-based calculation
         uint64 totalNonWithdrawnKeyCount;
         for (uint256 i = _startIndex; i < _endIndex; ) {
             uint256 validatorId = validatorIdsByOperatorId[operatorId][i];
@@ -570,7 +604,7 @@
             }
         }
         // If the result array isn't full, resize it to remove the unused elements
         assembly {
             mstore(validators, validatorCount)
         }
 
@@ -676,13 +710,21 @@
         validatorRegistry[_validatorId].status = ValidatorStatus.FRONT_RUN;
         uint256 operatorId = validatorRegistry[_validatorId].operatorId;
         operatorStructById[operatorId].active = false;
+        // decrement cached count if it was initialized earlier
+        if (operatorNonTerminalKeysCountInitialized[operatorId]) {
+            operatorNonTerminalKeysCount[operatorId]--;
+        }
     }
 
     // handle validator with invalid signature for 1ETH deposit
     //send back remaining ETH to operator address
     function handleInvalidSignature(uint256 _validatorId) internal {
         validatorRegistry[_validatorId].status = ValidatorStatus.INVALID_SIGNATURE;
         uint256 operatorId = validatorRegistry[_validatorId].operatorId;
+        // decrement cached count if it was initialized earlier
+        if (operatorNonTerminalKeysCountInitialized[operatorId]) {
+            operatorNonTerminalKeysCount[operatorId]--;
+        }
         address operatorAddress = operatorStructById[operatorId].operatorAddress;
         IOperatorRewardsCollector(staderConfig.getOperatorRewardsCollector()).depositFor{
             value: (COLLATERAL_ETH - staderConfig.getPreDepositSize())

contracts/interfaces/IPermissionlessNodeRegistry.sol

@@ -23,6 +23,7 @@ interface IPermissionlessNodeRegistry {
     event TransferredCollateralToPool(uint256 amount);
     event ValidatorAddedViaReferral(uint256 amount, string referralId);
     event UpdateMaxKeysPerOperator(uint256 maxKeysPerOperator);
+    event OperatorNonTerminalKeysCountSet(uint256 indexed operatorId, uint64 nonTerminalKeysCount);
 
     //Getters
 
@@ -82,6 +83,8 @@ interface IPermissionlessNodeRegistry {
         bool _optInForSocializingPool
     ) external returns (address mevFeeRecipientAddress);
 
+    function setOperatorNonTerminalKeysCount(uint256 _operatorId, uint64 _nonTerminalKeysCount) external;
+
     function pause() external;
 
     function unpause() external;