feat: poc for OperatorSets

src/ServiceManagerBase.sol

@@ -72,31 +72,37 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
     ) internal virtual onlyInitializing {
         _transferOwnership(initialOwner);
         _setRewardsInitiator(_rewardsInitiator);
-        isStrategiesMigrated = true;
     }
 
-    /// TODO: natspec
+    /**
+     * @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);
     }
 
-    /// TODO: natspec
+    /**
+     * @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 migrates all existing operators and strategies to operator sets
+    /**
+     * @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 {
-        require(
-            !isStrategiesMigrated,
-            "ServiceManagerBase.migrateStrategiesToOperatorSets: already migrated"
-        );
         uint8 quorumCount = _registryCoordinator.quorumCount();
         for (uint8 i = 0; i < quorumCount; ++i) {
             uint256 numStrategies = _stakeRegistry.strategyParamsLength(i);
@@ -105,24 +111,34 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
             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);
         }
-        isStrategiesMigrated = true;
     }
 
     /**
-     * @notice the operator needs to provide a signature over the operatorSetIds they will be registering
-     * for. This can be called externally by getOperatorSetIds
+     * @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(
-            !isOperatorMigrated[operator],
+            _operatorSetManager.avsOperatorStatus(address(this), operator)
+                == IOperatorSetManager.OperatorAVSRegistrationStatus.REGISTERED
+                && _operatorSetManager.operatorAVSOperatorSetCount(address(this), operator) == 0,
             "ServiceManagerBase.migrateOperatorToOperatorSets: already migrated"
         );
         uint32[] memory operatorSetIds = getOperatorSetIds(operator);
@@ -136,24 +152,24 @@ abstract contract ServiceManagerBase is OwnableUpgradeable, ServiceManagerBaseSt
      * 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 operatorSet The OperatorSet to eject the operators from
+     * @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,
-        IOperatorSetManager.OperatorSet calldata operatorSet
+        uint32 operatorSetId
     ) external onlyOwner {
         require(
-            operatorSet.avs == address(this),
-            "ServiceManagerBase.ejectNonmigratedOperators: operatorSet is not for this AVS"
-        );
-        require(
-            operatorSet.id < _registryCoordinator.quorumCount(),
+            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.isOperatorInOperatorSet(operators[i], operatorSet),
+                _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]);

src/ServiceManagerBaseStorage.sol

@@ -33,10 +33,6 @@ abstract contract ServiceManagerBaseStorage is IServiceManager {
     /// @notice The address of the entity that can initiate rewards
     address public rewardsInitiator;
 
-    bool internal isStrategiesMigrated;
-
-    mapping(address => bool) isOperatorMigrated;
-
     /// @notice Sets the (immutable) `_avsDirectory`, `_rewardsCoordinator`, `_registryCoordinator`, and `_stakeRegistry` addresses
     constructor(
         IOperatorSetManager __operatorSetManager,
@@ -51,5 +47,5 @@ abstract contract ServiceManagerBaseStorage is IServiceManager {
     }
 
     // storage gap for upgradeability
-    uint256[47] private __GAP;
+    uint256[49] private __GAP;
 }

src/StakeRegistry.sol

@@ -42,10 +42,9 @@ contract StakeRegistry is StakeRegistryStorage {
     }
 
     constructor(
-        IServiceManager _serviceManager,
         IRegistryCoordinator _registryCoordinator,
         IDelegationManager _delegationManager
-    ) StakeRegistryStorage(_serviceManager, _registryCoordinator, _delegationManager) {}
+    ) StakeRegistryStorage(_registryCoordinator, _delegationManager) {}
 
     /*******************************************************************************
                       EXTERNAL FUNCTIONS - REGISTRY COORDINATOR

src/StakeRegistryStorage.sol

@@ -50,15 +50,10 @@ abstract contract StakeRegistryStorage is IStakeRegistry {
 
 
     constructor(
-        IServiceManager _serviceManager,
         IRegistryCoordinator _registryCoordinator, 
         IDelegationManager _delegationManager
     ) {
-        require(
-            _registryCoordinator.serviceManager() == _serviceManager,
-            "StakeRegistryStorage: serviceManager does not match with registryCoordinator"
-        );
-        serviceManager = _serviceManager;
+        serviceManager = _registryCoordinator.serviceManager();
         registryCoordinator = address(_registryCoordinator);
         delegation = _delegationManager;
     }

src/interfaces/IOperatorSetManager.sol

@@ -149,6 +149,20 @@ interface IOperatorSetManager is ISignatureUtils {
 		OperatorSet calldata operatorSet
 	) external view returns(bool);
 
+    /**
+     * @notice retrieve the operator status for a given AVS, returns enum for either 
+     * DEREGISTERED or REGISTERED
+     * @param avs the AVS to get the operator status for
+     * @param operator the operator to get the status for
+     */
+    function avsOperatorStatus(address avs, address operator) external view returns (OperatorAVSRegistrationStatus);
+
+    /// @notice number of operatorSets the operator is registered for a given AVS
+    function operatorAVSOperatorSetCount(address avs, address operator) external view returns (uint256);
+
+    /// @notice mapping bool returning whether a strategy is part of an OperatorSet
+    function operatorSetStrategies(address avs, uint32 operatorSetId, IStrategy strategy) external view returns (bool);
+
     /**
      * @param operator the operator to get allowedToRegister for
      * @param operatorSet the operator set to get allowedToRegister for

src/interfaces/IServiceManager.sol

@@ -82,12 +82,12 @@ interface IServiceManager is IServiceManagerUI {
      * 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 operatorSet The OperatorSet to eject the operators from
+     * @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,
-        IOperatorSetManager.OperatorSet calldata operatorSet
+        uint32 operatorSetId
     ) external;
 
     /**

test/harnesses/StakeRegistryHarness.sol

@@ -8,10 +8,9 @@ import {IServiceManager} from "../../src/interfaces/IServiceManager.sol";
 // wrapper around the StakeRegistry contract that exposes the internal functions for unit testing.
 contract StakeRegistryHarness is StakeRegistry {
     constructor(
-        IServiceManager _serviceManager,
         IRegistryCoordinator _registryCoordinator,
         IDelegationManager _delegationManager
-    ) StakeRegistry(_serviceManager, _registryCoordinator, _delegationManager) {
+    ) StakeRegistry(_registryCoordinator, _delegationManager) {
     }
 
     function recordOperatorStakeUpdate(bytes32 operatorId, uint8 quorumNumber, uint96 newStake) external returns(int256) {

test/integration/IntegrationDeployer.t.sol

@@ -335,7 +335,7 @@ abstract contract IntegrationDeployer is Test, IUserDeployer {
         cheats.stopPrank();
 
         StakeRegistry stakeRegistryImplementation = new StakeRegistry(
-            serviceManager, IRegistryCoordinator(registryCoordinator), IDelegationManager(delegationManager)
+            IRegistryCoordinator(registryCoordinator), IDelegationManager(delegationManager)
         );
         BLSApkRegistry blsApkRegistryImplementation =
             new BLSApkRegistry(IRegistryCoordinator(registryCoordinator));

test/mocks/OperatorSetManagerMock.sol

@@ -95,7 +95,21 @@ contract OperatorSetManagerMock is IOperatorSetManager {
 		IStrategy[] calldata strategies
 	) external {}
 
-        /**
+    /**
+     * @notice retrieve the operator status for a given AVS, returns enum for either 
+     * DEREGISTERED or REGISTERED
+     * @param avs the AVS to get the operator status for
+     * @param operator the operator to get the status for
+     */
+    function avsOperatorStatus(address avs, address operator) external view returns (OperatorAVSRegistrationStatus) {}
+
+    /// @notice number of operatorSets the operator is registered for a given AVS
+    function operatorAVSOperatorSetCount(address avs, address operator) external view returns (uint256) {}
+
+    /// @notice mapping bool returning whether a strategy is part of an OperatorSet
+    function operatorSetStrategies(address avs, uint32 operatorSetId, IStrategy strategy) external view returns (bool) {}
+
+    /**
      * @param operator the operator to get allowedToRegister for
      * @param operatorSet the operator set to get allowedToRegister for
      *

test/unit/StakeRegistryUnit.t.sol

@@ -52,7 +52,6 @@ contract StakeRegistryUnitTests is MockAVSDeployer, IStakeRegistryEvents {
         );
 
         stakeRegistryImplementation = new StakeRegistryHarness(
-            serviceManager,
             IRegistryCoordinator(address(registryCoordinator)),
             delegationMock
         );

test/utils/MockAVSDeployer.sol

@@ -225,7 +225,7 @@ contract MockAVSDeployer is Test {
         cheats.startPrank(proxyAdminOwner);
 
         stakeRegistryImplementation =
-            new StakeRegistryHarness(serviceManager, IRegistryCoordinator(registryCoordinator), delegationMock);
+            new StakeRegistryHarness(IRegistryCoordinator(registryCoordinator), delegationMock);
 
         proxyAdmin.upgrade(
             TransparentUpgradeableProxy(payable(address(stakeRegistry))),