feat: poc for OperatorSets

src/RegistryCoordinator.sol

@@ -495,18 +495,15 @@ contract RegistryCoordinator is
         emit OperatorSocketUpdate(operatorId, socket);
 
         // If the operator wasn't registered for any quorums, update their status
-        // and register them with this AVS in EigenLayer core (DelegationManager)
         if (_operatorInfo[operator].status != OperatorStatus.REGISTERED) {
             _operatorInfo[operator] = OperatorInfo({
                 operatorId: operatorId,
                 status: OperatorStatus.REGISTERED
             });
-
-            // Register the operator with the EigenLayer core contracts via this AVS's ServiceManager
-            serviceManager.registerOperatorToAVS(operator, operatorSignature);
-
-            emit OperatorRegistered(operator, operatorId);
         }
+        // Register the operator with the EigenLayer core contracts via this AVS's ServiceManager
+        serviceManager.registerOperatorToOperatorSets(operator, quorumNumbers, operatorSignature);
+        emit OperatorRegistered(operator, operatorId, quorumNumbers);
 
         // Register the operator with the BLSApkRegistry, StakeRegistry, and IndexRegistry
         blsApkRegistry.registerOperator(operator, quorumNumbers);
@@ -612,13 +609,12 @@ contract RegistryCoordinator is
             newBitmap: newBitmap
         });
 
-        // If the operator is no longer registered for any quorums, update their status and deregister 
-        // them from the AVS via the EigenLayer core contracts
+        // If the operator is no longer registered for any quorums, update their status
         if (newBitmap.isEmpty()) {
             operatorInfo.status = OperatorStatus.DEREGISTERED;
-            serviceManager.deregisterOperatorFromAVS(operator);
-            emit OperatorDeregistered(operator, operatorId);
         }
+        serviceManager.deregisterOperatorFromOperatorSets(operator, quorumNumbers);
+        emit OperatorDeregistered(operator, operatorId, quorumNumbers);
 
         // Deregister operator with each of the registry contracts
         blsApkRegistry.deregisterOperator(operator, quorumNumbers);

src/ServiceManagerBase.sol

@@ -4,9 +4,10 @@ pragma solidity ^0.8.12;
 import {OwnableUpgradeable} from "@openzeppelin-upgrades/contracts/access/OwnableUpgradeable.sol";
 import {Initializable} from "@openzeppelin-upgrades/contracts/proxy/utils/Initializable.sol";
 import {ISignatureUtils} from "eigenlayer-contracts/src/contracts/interfaces/ISignatureUtils.sol";
-import {IAVSDirectory} from "eigenlayer-contracts/src/contracts/interfaces/IAVSDirectory.sol";
-import {IRewardsCoordinator} from "eigenlayer-contracts/src/contracts/interfaces/IRewardsCoordinator.sol";
+import {IRewardsCoordinator} from
+    "eigenlayer-contracts/src/contracts/interfaces/IRewardsCoordinator.sol";
 
+import {IOperatorSetManager, IStrategy} from "./interfaces/IOperatorSetManager.sol"; // should be later changed to be import from core
 import {ServiceManagerBaseStorage} from "./ServiceManagerBaseStorage.sol";
 import {IServiceManager} from "./interfaces/IServiceManager.sol";
 import {IRegistryCoordinator} from "./interfaces/IRegistryCoordinator.sol";
@@ -30,6 +31,15 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
         _;
     }
 
+    /// @notice when applied to a function, only allows the StakeRegistry to call it
+    modifier onlyStakeRegistry() {
+        require(
+            msg.sender == address(_stakeRegistry),
+            "ServiceManagerBase.onlyStakeRegistry: caller is not the stake registry"
+        );
+        _;
+    }
+
     /// @notice only rewardsInitiator can call createAVSRewardsSubmission
     modifier onlyRewardsInitiator() {
         require(
@@ -41,13 +51,13 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
 
     /// @notice Sets the (immutable) `_registryCoordinator` address
     constructor(
-        IAVSDirectory __avsDirectory,
+        IOperatorSetManager __operatorSetManager,
         IRewardsCoordinator __rewardsCoordinator,
         IRegistryCoordinator __registryCoordinator,
         IStakeRegistry __stakeRegistry
     )
         ServiceManagerBaseStorage(
-            __avsDirectory,
+            __operatorSetManager,
             __rewardsCoordinator,
             __registryCoordinator,
             __stakeRegistry
@@ -64,13 +74,118 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
         _setRewardsInitiator(_rewardsInitiator);
     }
 
+    /**
+     * @notice called by the AVS StakeRegistry whenever a new IStrategy is added to a quorum/operatorSet
+     * @dev calls operatorSetManager.addStrategiesToOperatorSet()
+     */
+    function addStrategiesToOperatorSet(
+        uint32 operatorSetID,
+        IStrategy[] calldata strategies
+    ) external onlyStakeRegistry {
+        _operatorSetManager.addStrategiesToOperatorSet(operatorSetID, strategies);
+    }
+
+    /**
+     * @notice called by the AVS StakeRegistry whenever a new IStrategy is removed from a quorum/operatorSet
+     * @dev calls operatorSetManager.removeStrategiesFromOperatorSet()
+     */
+    function removeStrategiesFromOperatorSet(
+        uint32 operatorSetID,
+        IStrategy[] calldata strategies
+    ) external onlyStakeRegistry {
+        _operatorSetManager.removeStrategiesFromOperatorSet(operatorSetID, strategies);
+    }
+
+    /**
+     * @notice One-time call that migrates all existing strategies for each quorum to their respective operator sets
+     * Note: a separate migration per operator must be performed to migrate an existing operator to the operator set
+     * See migrateOperatorToOperatorSets() below
+     * @dev calls operatorSetManager.addStrategiesToOperatorSet()
+     */
+    function migrateStrategiesToOperatorSets() external {
+        uint8 quorumCount = _registryCoordinator.quorumCount();
+        for (uint8 i = 0; i < quorumCount; ++i) {
+            uint256 numStrategies = _stakeRegistry.strategyParamsLength(i);
+            IStrategy[] memory strategies = new IStrategy[](numStrategies);
+            // get the strategies for the quorum/operator set
+            for (uint256 j = 0; j < numStrategies; ++j) {
+                IStrategy strategy = _stakeRegistry.strategyParamsByIndex(i, j).strategy;
+                strategies[j] = strategy;
+                require(
+                    !_operatorSetManager.operatorSetStrategies(address(this), uint32(i), strategy),
+                    "ServiceManagerBase.migrateStrategiesToOperatorSets: strategy already part of OperatorSet"
+                );
+            }
+
+            _operatorSetManager.addStrategiesToOperatorSet(uint32(i), strategies);
+            emit OperatorSetStrategiesMigrated(uint32(i), strategies);
+        }
+    }
+
+    /**
+     * @notice One-time call to migrate an existing operator to the respective operator sets.
+     * The operator needs to provide a signature over the operatorSetIds they are currently registered for.
+     * This can be retrieved externally by calling getOperatorSetIds.
+     * @param operator the address of the operator to be migrated
+     * @param signature the signature of the operator on their intent to migrate
+     * @dev calls operatorSetManager.registerOperatorToOperatorSets()
+     */
+    function migrateOperatorToOperatorSets(
+        address operator,
+        ISignatureUtils.SignatureWithSaltAndExpiry memory signature
+    ) external {
+        // check if operator is already migrated, that is they are registered operators with 0 operatorSet count
+        require(
+            _operatorSetManager.avsOperatorStatus(address(this), operator)
+                == IOperatorSetManager.OperatorAVSRegistrationStatus.REGISTERED
+                && _operatorSetManager.operatorAVSOperatorSetCount(address(this), operator) == 0,
+            "ServiceManagerBase.migrateOperatorToOperatorSets: already migrated"
+        );
+        uint32[] memory operatorSetIds = getOperatorSetIds(operator);
+
+        _operatorSetManager.registerOperatorToOperatorSets(operator, operatorSetIds, signature);
+        emit OperatorMigratedToOperatorSets(operator, operatorSetIds);
+    }
+
+    /**
+     * @notice Once strategies have been migrated and operators have been migrated to operator sets.
+     * The ServiceManager owner can eject any operators that have yet to completely migrate fully to operator sets.
+     * This final step of the migration process will ensure the full migration of all operators to operator sets.
+     * @param operators The list of operators to eject for the given OperatorSet
+     * @param operatorSetId This AVS's operatorSetId to eject operators from
+     * @dev The RegistryCoordinator MUST set this ServiceManager contract to be the ejector address for this call to succeed
+     */
+    function ejectNonmigratedOperators(
+        address[] calldata operators,
+        uint32 operatorSetId
+    ) external onlyOwner {
+        require(
+            operatorSetId < _registryCoordinator.quorumCount(),
+            "ServiceManagerBase.ejectNonmigratedOperators: operatorSet does not exist"
+        );
+        IOperatorSetManager.OperatorSet memory operatorSet =
+            IOperatorSetManager.OperatorSet({avs: address(this), id: operatorSetId});
+        for (uint256 i = 0; i < operators.length; ++i) {
+            require(
+                _operatorSetManager.avsOperatorStatus(address(this), operator)
+                    == IOperatorSetManager.OperatorAVSRegistrationStatus.REGISTERED
+                    && !_operatorSetManager.isOperatorInOperatorSet(operators[i], operatorSet),
+                "ServiceManagerBase.removeNonmigratedOperators: operator already registered to operator set"
+            );
+            bytes32 operatorId = _registryCoordinator.getOperatorId(operators[i]);
+            uint192 quorumBitmap = _registryCoordinator.getCurrentQuorumBitmap(operatorId);
+            bytes memory quorumNumbers = BitmapUtils.bitmapToBytesArray(quorumBitmap);
+            _registryCoordinator.ejectOperator(operators[i], quorumNumbers);
+        }
+    }
+
     /**
      * @notice Updates the metadata URI for the AVS
      * @param _metadataURI is the metadata URI for the AVS
      * @dev only callable by the owner
      */
     function updateAVSMetadataURI(string memory _metadataURI) public virtual onlyOwner {
-        _avsDirectory.updateAVSMetadataURI(_metadataURI);
+        _operatorSetManager.updateAVSMetadataURI(_metadataURI);
     }
 
     /**
@@ -84,15 +199,15 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
      * @dev This function will revert if the `rewardsSubmission` is malformed,
      * e.g. if the `strategies` and `weights` arrays are of non-equal lengths
      */
-    function createAVSRewardsSubmission(IRewardsCoordinator.RewardsSubmission[] calldata rewardsSubmissions)
-        public
-        virtual
-        onlyRewardsInitiator
-    {
+    function createAVSRewardsSubmission(
+        IRewardsCoordinator.RewardsSubmission[] calldata rewardsSubmissions
+    ) public virtual onlyRewardsInitiator {
         for (uint256 i = 0; i < rewardsSubmissions.length; ++i) {
             // transfer token to ServiceManager and approve RewardsCoordinator to transfer again
             // in createAVSRewardsSubmission() call
-            rewardsSubmissions[i].token.transferFrom(msg.sender, address(this), rewardsSubmissions[i].amount);
+            rewardsSubmissions[i].token.transferFrom(
+                msg.sender, address(this), rewardsSubmissions[i].amount
+            );
             uint256 allowance =
                 rewardsSubmissions[i].token.allowance(address(this), address(_rewardsCoordinator));
             rewardsSubmissions[i].token.approve(
@@ -112,15 +227,50 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
         address operator,
         ISignatureUtils.SignatureWithSaltAndExpiry memory operatorSignature
     ) public virtual onlyRegistryCoordinator {
-        _avsDirectory.registerOperatorToAVS(operator, operatorSignature);
+        _operatorSetManager.registerOperatorToAVS(operator, operatorSignature);
+    }
+
+    /**
+     * @notice Called by this AVS's RegistryCoordinator to register an operator for its registering operatorSets
+     *
+     * @param operator the address of the operator to be added to the operator set
+     * @param quorumNumbers quorums/operatorSetIds to add the operator to
+     * @param signature the signature of the operator on their intent to register
+     * @dev msg.sender is used as the AVS
+     * @dev operator must not have a pending a deregistration from the operator set
+     * @dev if this is the first operator set in the AVS that the operator is
+     * registering for, a OperatorAVSRegistrationStatusUpdated event is emitted with
+     * a REGISTERED status
+     */
+    function registerOperatorToOperatorSets(
+        address operator,
+        bytes calldata quorumNumbers,
+        ISignatureUtils.SignatureWithSaltAndExpiry memory signature
+    ) external onlyRegistryCoordinator {
+        uint32[] memory operatorSetIds = quorumsToOperatorSetIds(quorumNumbers);
+        _operatorSetManager.registerOperatorToOperatorSets(operator, operatorSetIds, signature);
     }
 
     /**
-     * @notice Forwards a call to EigenLayer's AVSDirectory contract to confirm operator deregistration from the AVS
-     * @param operator The address of the operator to deregister.
+     * @notice Called by this AVS's RegistryCoordinator to deregister an operator for its operatorSets
+     *
+     * @param operator the address of the operator to be removed from the
+     * operator set
+     * @param quorumNumbers the quorumNumbers/operatorSetIds to deregister the operator for
+     *
+     * @dev msg.sender is used as the AVS
+     * @dev operator must be registered for msg.sender AVS and the given
+     * operator set
+     * @dev if this removes operator from all operator sets for the msg.sender AVS
+     * then an OperatorAVSRegistrationStatusUpdated event is emitted with a DEREGISTERED
+     * status
      */
-    function deregisterOperatorFromAVS(address operator) public virtual onlyRegistryCoordinator {
-        _avsDirectory.deregisterOperatorFromAVS(operator);
+    function deregisterOperatorFromOperatorSets(
+        address operator,
+        bytes calldata quorumNumbers
+    ) external onlyRegistryCoordinator {
+        uint32[] memory operatorSetIds = quorumsToOperatorSetIds(quorumNumbers);
+        _operatorSetManager.deregisterOperatorFromOperatorSets(operator, operatorSetIds);
     }
 
     /**
@@ -209,8 +359,30 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
         return restakedStrategies;
     }
 
-    /// @notice Returns the EigenLayer AVSDirectory contract.
-    function avsDirectory() external view override returns (address) {
-        return address(_avsDirectory);
+    /// @notice converts the quorumNumbers array to operatorSetIds array
+    function quorumsToOperatorSetIds(bytes memory quorumNumbers)
+        internal
+        pure
+        returns (uint32[] memory)
+    {
+        uint256 quorumCount = quorumNumbers.length;
+        uint32[] memory operatorSetIds = new uint32[](quorumCount);
+        for (uint256 i = 0; i < quorumCount; i++) {
+            operatorSetIds[i] = uint32(uint8(quorumNumbers[i]));
+        }
+        return operatorSetIds;
+    }
+
+    /// @notice Returns the operator set IDs for the given operator address by querying the RegistryCoordinator
+    function getOperatorSetIds(address operator) public view returns (uint32[] memory) {
+        bytes32 operatorId = _registryCoordinator.getOperatorId(operator);
+        uint192 operatorBitmap = _registryCoordinator.getCurrentQuorumBitmap(operatorId);
+        bytes memory operatorQuorums = BitmapUtils.bitmapToBytesArray(operatorBitmap);
+        return quorumsToOperatorSetIds(operatorQuorums);
+    }
+
+    /// @notice Returns the EigenLayer OperatorSetManager contract.
+    function operatorSetManager() external view override returns (address) {
+        return address(_operatorSetManager);
     }
 }

src/ServiceManagerBaseStorage.sol

@@ -5,7 +5,7 @@ import {IServiceManager} from "./interfaces/IServiceManager.sol";
 import {IRegistryCoordinator} from "./interfaces/IRegistryCoordinator.sol";
 import {IStakeRegistry} from "./interfaces/IStakeRegistry.sol";
 
-import {IAVSDirectory} from "eigenlayer-contracts/src/contracts/interfaces/IAVSDirectory.sol";
+import {IOperatorSetManager} from "./interfaces/IOperatorSetManager.sol"; // should be later changed to be import from core
 import {IRewardsCoordinator} from "eigenlayer-contracts/src/contracts/interfaces/IRewardsCoordinator.sol";
 
 /**
@@ -19,7 +19,7 @@ abstract contract ServiceManagerBaseStorage is IServiceManager {
      *                            CONSTANTS AND IMMUTABLES
      *
      */
-    IAVSDirectory internal immutable _avsDirectory;
+    IOperatorSetManager internal immutable _operatorSetManager;
     IRewardsCoordinator internal immutable _rewardsCoordinator;
     IRegistryCoordinator internal immutable _registryCoordinator;
     IStakeRegistry internal immutable _stakeRegistry;
@@ -35,12 +35,12 @@ abstract contract ServiceManagerBaseStorage is IServiceManager {
 
     /// @notice Sets the (immutable) `_avsDirectory`, `_rewardsCoordinator`, `_registryCoordinator`, and `_stakeRegistry` addresses
     constructor(
-        IAVSDirectory __avsDirectory,
+        IOperatorSetManager __operatorSetManager,
         IRewardsCoordinator __rewardsCoordinator,
         IRegistryCoordinator __registryCoordinator,
         IStakeRegistry __stakeRegistry
     ) {
-        _avsDirectory = __avsDirectory;
+        _operatorSetManager = __operatorSetManager;
         _rewardsCoordinator = __rewardsCoordinator;
         _registryCoordinator = __registryCoordinator;
         _stakeRegistry = __stakeRegistry;