Use uint32 for number of ejected operators (#232)

* Use uint32 for number of ejected operators * fix test
Layr-Labs · Apr 17, 2024 · 59ebd32 · 59ebd32
1 parent 0328639
commit 59ebd32
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 37 deletions.
diff --git a/src/EjectionManager.sol b/src/EjectionManager.sol
@@ -23,13 +23,13 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
     /// @notice Addresses permissioned to eject operators under a ratelimit
     mapping(address => bool) public isEjector;
 
-    /// @notice Keeps track of the total stake ejected for a quorum 
+    /// @notice Keeps track of the total stake ejected for a quorum
     mapping(uint8 => StakeEjection[]) public stakeEjectedForQuorum;
     /// @notice Ratelimit parameters for each quorum
     mapping(uint8 => QuorumEjectionParams) public quorumEjectionParams;
 
     constructor(
-        IRegistryCoordinator _registryCoordinator, 
+        IRegistryCoordinator _registryCoordinator,
         IStakeRegistry _stakeRegistry
     ) {
         registryCoordinator = _registryCoordinator;
@@ -44,7 +44,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
      * @param _quorumEjectionParams are the ratelimit parameters for the quorum at each index
      */
     function initialize(
-        address _owner, 
+        address _owner,
         address[] memory _ejectors,
         QuorumEjectionParams[] memory _quorumEjectionParams
     ) external initializer {
@@ -64,25 +64,25 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
      * @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 returns (uint256[] memory){
+    function ejectOperators(bytes32[][] memory _operatorIds) external returns (uint32[] memory){
         require(isEjector[msg.sender] || msg.sender == owner(), "Ejector: Only owner or ejector can eject");
 
-        uint256[] memory ejectedOperatorsForQuorum = new uint256[](_operatorIds.length);
+        uint32[] memory ejectedOperatorsForQuorum = new uint32[](_operatorIds.length);
 
         for(uint i = 0; i < _operatorIds.length; ++i) {
             uint8 quorumNumber = uint8(i);
 
             uint256 amountEjectable = amountEjectableForQuorum(quorumNumber);
             uint256 stakeForEjection;
-            uint256 ejectedOperators;
+            uint32 ejectedOperators;
 
             bool broke;
             for(uint8 j = 0; j < _operatorIds[i].length; ++j) {
                 uint256 operatorStake = stakeRegistry.getCurrentStake(_operatorIds[i][j], quorumNumber);
 
                 //if caller is ejector enforce ratelimit
                 if(
-                    isEjector[msg.sender] && 
+                    isEjector[msg.sender] &&
                     quorumEjectionParams[quorumNumber].rateLimitWindow > 0 &&
                     stakeForEjection + operatorStake > amountEjectable
                 ){
@@ -108,7 +108,7 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
             }
 
             //record the stake ejected if ejector and ratelimit enforced
-            if(!broke && isEjector[msg.sender]){ 
+            if(!broke && isEjector[msg.sender]){
                 stakeEjectedForQuorum[quorumNumber].push(StakeEjection({
                     timestamp: block.timestamp,
                     stakeEjected: stakeForEjection
@@ -177,6 +177,6 @@ contract EjectionManager is IEjectionManager, OwnableUpgradeable{
         if(totalEjected >= totalEjectable){
             return 0;
         }
-        return totalEjectable - totalEjected;   
+        return totalEjectable - totalEjected;
     }
-}
+}
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 returns (uint256[] memory);
+    function ejectOperators(bytes32[][] memory _operatorIds) external returns (uint32[] memory);
 
     /**
      * @notice Sets the ratelimit parameters for a quorum
@@ -52,4 +52,4 @@ interface IEjectionManager {
      * @param _quorumNumber The quorum number to view ejectable stake for
      */
     function amountEjectableForQuorum(uint8 _quorumNumber) external view returns (uint256);
-}
+}
diff --git a/test/unit/EjectionManagerUnit.t.sol b/test/unit/EjectionManagerUnit.t.sol
@@ -77,7 +77,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
             }
         }
 
-        assertEq(uint8(registryCoordinator.getOperatorStatus(defaultOperator)), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+        assertEq(uint8(registryCoordinator.getOperatorStatus(defaultOperator)), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
 
         for(uint8 i = 0; i < numQuorums; i++) {
             for(uint8 j = 0; j < operatorsToEject; j++) {
@@ -87,7 +87,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+        uint32[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
 
         for(uint8 i = 0; i < numQuorums; i++) {
             assertEq(ejectedOperatorsForQuorum[i], operatorsToEject);
@@ -111,7 +111,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
 
         for(uint8 i = 0; i < numQuorums; i++) {
@@ -122,14 +122,14 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+        uint32[] 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)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
         }
     }
 
@@ -149,7 +149,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
 
         for(uint8 i = 0; i < numQuorums; i++) {
@@ -160,18 +160,18 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+        uint32[] 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)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
         }
 
         for(uint8 i = operatorsCanEject; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
     }
 
@@ -190,7 +190,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
 
         for(uint8 i = 0; i < numQuorums; i++) {
@@ -201,14 +201,14 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+        uint32[] 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)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
         }
 
         cheats.warp(block.timestamp + (ratelimitWindow / 2));
@@ -222,7 +222,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, operatorsToEject + i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, operatorsToEject + i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
 
         for(uint8 i = 0; i < numQuorums; i++) {
@@ -240,7 +240,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, operatorsToEject + i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, operatorsToEject + i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
         }
     }
 
@@ -264,7 +264,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
 
         for(uint8 i = 0; i < numQuorums; i++) {
@@ -275,14 +275,14 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+        uint32[] 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)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
         }
     }
 
@@ -301,7 +301,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         for(uint8 i = 0; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
 
         for(uint8 i = 0; i < numQuorums; i++) {
@@ -312,14 +312,14 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(registryCoordinatorOwner);
-        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+        uint32[] 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)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
         }
     }
 
@@ -341,7 +341,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         registryCoordinator.deregisterOperator(BitmapUtils.bitmapToBytesArray(MAX_QUORUM_BITMAP));
 
         for(uint8 i = 1; i < operatorsToEject; i++) {
-            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.REGISTERED));
         }
 
         for(uint8 i = 0; i < numQuorums; i++) {
@@ -354,14 +354,14 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         }
 
         cheats.prank(ejector);
-        uint256[] memory ejectedOperatorsForQuorum = ejectionManager.ejectOperators(operatorIds);
+        uint32[] 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)); 
+            assertEq(uint8(registryCoordinator.getOperatorStatus(_incrementAddress(defaultOperator, i))), uint8(IRegistryCoordinator.OperatorStatus.DEREGISTERED));
         }
     }
 
@@ -412,7 +412,7 @@ contract EjectionManagerUnitTests is MockAVSDeployer {
         for (uint i = 0; i < numOperators; i++) {
             BN254.G1Point memory pubKey = BN254.hashToG1(keccak256(abi.encodePacked(i)));
             address operator = _incrementAddress(defaultOperator, i);
-            _registerOperatorWithCoordinator(operator, MAX_QUORUM_BITMAP, pubKey, stake); 
+            _registerOperatorWithCoordinator(operator, MAX_QUORUM_BITMAP, pubKey, stake);
         }
     }
-}
+}