fix: correctly set bit in stake registry (#112)

src/BLSSignatureChecker.sol

@@ -219,7 +219,7 @@ contract BLSSignatureChecker is IBLSSignatureChecker {
                 // if so, load their stake at referenceBlockNumber and subtract it from running stake signed
                 for (uint256 j = 0; j < params.nonSignerPubkeys.length; j++) {
                     // if the nonSigner is a part of the quorum, subtract their stake from the running total
-                    if (BitmapUtils.numberIsInBitmap(nonSigners.quorumBitmaps[j], uint8(quorumNumbers[i]))) {
+                    if (BitmapUtils.isSet(nonSigners.quorumBitmaps[j], uint8(quorumNumbers[i]))) {
                         stakeTotals.signedStakeForQuorum[i] -=
                             stakeRegistry.getStakeAtBlockNumberAndIndex({
                                 quorumNumber: uint8(quorumNumbers[i]),

src/RegistryCoordinator.sol

@@ -331,7 +331,7 @@ contract RegistryCoordinator is
                 {
                     uint192 currentBitmap = _currentOperatorBitmap(operatorId);
                     require(
-                        BitmapUtils.numberIsInBitmap(currentBitmap, quorumNumber),
+                        BitmapUtils.isSet(currentBitmap, quorumNumber),
                         "RegistryCoordinator.updateOperatorsForQuorum: operator not in quorum"
                     );
                     // Require check is to prevent duplicate operators and that all quorum operators are updated

src/StakeRegistry.sol

@@ -20,6 +20,8 @@ import {BitmapUtils} from "src/libraries/BitmapUtils.sol";
  * @author Layr Labs, Inc.
  */
 contract StakeRegistry is StakeRegistryStorage {
+
+    using BitmapUtils for *;
 
     modifier onlyRegistryCoordinator() {
         require(
@@ -138,7 +140,7 @@ contract StakeRegistry is StakeRegistryStorage {
      * or more quorums.
      *
      * If the operator no longer has the minimum stake required for a quorum, they are
-     * added to the
+     * added to the `quorumsToRemove`, which is returned to the registry coordinator
      * @return A bitmap of quorums where the operator no longer meets the minimum stake
      * and should be deregistered.
      */
@@ -168,7 +170,7 @@ contract StakeRegistry is StakeRegistryStorage {
             // If the operator no longer meets the minimum stake, set their stake to zero and mark them for removal
             if (!hasMinimumStake) {
                 stakeWeight = 0;
-                quorumsToRemove |= quorumNumber;
+                quorumsToRemove = uint192(quorumsToRemove.setBit(quorumNumber));
             }
 
             // Update the operator's stake and retrieve the delta

src/libraries/BitmapUtils.sol

@@ -295,9 +295,19 @@ library BitmapUtils {
         return count;
     }
 
-    /// @notice returns 'true' if `numberToCheckForInclusion` is in `bitmap` and 'false' otherwise.
-    function numberIsInBitmap(uint256 bitmap, uint8 numberToCheckForInclusion) internal pure returns (bool) {
-        return (((bitmap >> numberToCheckForInclusion) & 1) == 1);
+    /// @notice Returns `true` if `bit` is in `bitmap`. Returns `false` otherwise.
+    function isSet(uint256 bitmap, uint8 bit) internal pure returns (bool) {
+        return 1 == ((bitmap >> bit) & 1);
+    }
+
+    /**
+     * @notice Returns a copy of `bitmap` with `bit` set. 
+     * @dev IMPORTANT: we're dealing with stack values here, so this doesn't modify
+     * the original bitmap. Using this correctly requires an assignment statement:
+     * `bitmap = bitmap.setBit(bit);`
+     */
+    function setBit(uint256 bitmap, uint8 bit) internal pure returns (uint256) {
+        return bitmap | (1 << bit);
     }
 
     /**

test/harnesses/BitmapUtilsWrapper.sol

@@ -33,7 +33,7 @@ contract BitmapUtilsWrapper {
         return BitmapUtils.countNumOnes(n);
     }
 
-    function numberIsInBitmap(uint256 bitmap, uint8 numberToCheckForInclusion) external pure returns (bool) {
-        return BitmapUtils.numberIsInBitmap(bitmap, numberToCheckForInclusion);
+    function isSet(uint256 bitmap, uint8 numberToCheckForInclusion) external pure returns (bool) {
+        return BitmapUtils.isSet(bitmap, numberToCheckForInclusion);
     }
 }

test/unit/BitmapUtils.t.sol

@@ -166,15 +166,15 @@ contract BitmapUtilsUnitTests is Test {
         require(libraryOutput == numOnes, "inconsistency in countNumOnes function");
     }
 
-    // @notice some simple sanity checks on the `numberIsInBitmap` function
+    // @notice some simple sanity checks on the `isSet` function
     function testNumberIsInBitmap() public view {
-        require(bitmapUtilsWrapper.numberIsInBitmap(2 ** 6, 6), "numberIsInBitmap function is broken 0");
-        require(bitmapUtilsWrapper.numberIsInBitmap(1, 0), "numberIsInBitmap function is broken 1");
-        require(bitmapUtilsWrapper.numberIsInBitmap(255, 7), "numberIsInBitmap function is broken 2");
-        require(bitmapUtilsWrapper.numberIsInBitmap(1024, 10), "numberIsInBitmap function is broken 3");
+        require(bitmapUtilsWrapper.isSet(2 ** 6, 6), "isSet function is broken 0");
+        require(bitmapUtilsWrapper.isSet(1, 0), "isSet function is broken 1");
+        require(bitmapUtilsWrapper.isSet(255, 7), "isSet function is broken 2");
+        require(bitmapUtilsWrapper.isSet(1024, 10), "isSet function is broken 3");
         for (uint256 i = 0; i < 256; ++i) {
-            require(bitmapUtilsWrapper.numberIsInBitmap(type(uint256).max, uint8(i)), "numberIsInBitmap function is broken 4");
-            require(!bitmapUtilsWrapper.numberIsInBitmap(0, uint8(i)), "numberIsInBitmap function is broken 5");
+            require(bitmapUtilsWrapper.isSet(type(uint256).max, uint8(i)), "isSet function is broken 4");
+            require(!bitmapUtilsWrapper.isSet(0, uint8(i)), "isSet function is broken 5");
         }
     }
 }