feat: return ops ejected (#228)

Layr-Labs · Apr 16, 2024 · 0328639 · 0328639
1 parent 6454c05
commit 0328639
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 11 deletions.
diff --git a/src/EjectionManager.sol b/src/EjectionManager.sol
@@ -60,17 +60,21 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
     /**
      * @notice Ejects operators from the AVSs RegistryCoordinator under a ratelimit
      * @param _operatorIds The ids of the operators 'j' to eject for each quorum 'i'
+     * @return ejectedOperatorsForQuorum The total number of operators ejected for each quorum 'i'
      * @dev This function will eject as many operators as possible without reverting prioritizing operators at the lower index
      * @dev The owner can eject operators without recording of stake ejection
      */
-    function ejectOperators(bytes32[][] memory _operatorIds) external {
+    function ejectOperators(bytes32[][] memory _operatorIds) external returns (uint256[] memory){
         require(isEjector[msg.sender] || msg.sender == owner(), "Ejector: Only owner or ejector can eject");
 
+        uint256[] memory ejectedOperatorsForQuorum = new uint256[](_operatorIds.length);
+
         for(uint i = 0; i < _operatorIds.length; ++i) {
             uint8 quorumNumber = uint8(i);
 
             uint256 amountEjectable = amountEjectableForQuorum(quorumNumber);
-            uint256 stakeForEjection; 
+            uint256 stakeForEjection;
+            uint256 ejectedOperators;
 
             bool broke;
             for(uint8 j = 0; j < _operatorIds[i].length; ++j) {
@@ -96,6 +100,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
                     abi.encodePacked(quorumNumber)
                 ) {
                     stakeForEjection += operatorStake;
+                    ++ejectedOperators;
                     emit OperatorEjected(_operatorIds[i][j], quorumNumber);
                 } catch (bytes memory err) {
                     emit FailedOperatorEjection(_operatorIds[i][j], quorumNumber, err);
@@ -110,7 +115,10 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
                 }));
             }
 
+            ejectedOperatorsForQuorum[i] = ejectedOperators;
         }
+
+        return ejectedOperatorsForQuorum;
     }
 
     /**

diff --git a/src/interfaces/IEjectionManager.sol b/src/interfaces/IEjectionManager.sol
@@ -32,7 +32,7 @@ interface IEjectionManager {
      * @notice Ejects operators from the AVSs registryCoordinator under a ratelimit
      * @param _operatorIds The ids of the operators to eject for each quorum
      */
-    function ejectOperators(bytes32[][] memory _operatorIds) external;
+    function ejectOperators(bytes32[][] memory _operatorIds) external returns (uint256[] memory);
 
     /**
      * @notice Sets the ratelimit parameters for a quorum

diff --git a/test/unit/EjectionManagerUnit.t.sol b/test/unit/EjectionManagerUnit.t.sol
@@ -87,7 +87,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        ejectionManager.ejectOperators(operatorIds);
+        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsToEject);
+        }
 
         assertEq(uint8(registryCoordinator.getOperatorStatus(defaultOperator)), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
     }
@@ -118,7 +122,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        ejectionManager.ejectOperators(operatorIds);
+        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsToEject);
+        }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
             assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); 
@@ -152,7 +160,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        ejectionManager.ejectOperators(operatorIds);
+        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsCanEject);
+        }
 
         for(uint8 i = 0; i < operatorsCanEject; i++) {
             assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); 
@@ -189,7 +201,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        ejectionManager.ejectOperators(operatorIds);
+        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsToEject);
+        }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
             assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); 
@@ -217,7 +233,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        ejectionManager.ejectOperators(operatorIds);
+        ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsToEject);
+        }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
             assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, operatorsToEject + i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); 
@@ -255,7 +275,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        ejectionManager.ejectOperators(operatorIds);
+        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsToEject);
+        }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
             assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); 
@@ -288,7 +312,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(registryCoordinatorOwner);
-        ejectionManager.ejectOperators(operatorIds);
+        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsToEject);
+        }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
             assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); 
@@ -326,7 +354,11 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        ejectionManager.ejectOperators(operatorIds);
+        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+
+        for(uint8 i = 0; i < numQuorums; i++) {
+            assertEq(ejectedOperatorsForQuorum[i], operatorsToEject - 1);
+        }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
             assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));