diff --git a/src/BLSSignatureChecker.sol b/src/BLSSignatureChecker.sol index 99d501ba..03dbe42d 100644 --- a/src/BLSSignatureChecker.sol +++ b/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]), diff --git a/src/RegistryCoordinator.sol b/src/RegistryCoordinator.sol index c3ec5aa6..9345ac5b 100644 --- a/src/RegistryCoordinator.sol +++ b/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 diff --git a/src/StakeRegistry.sol b/src/StakeRegistry.sol index 5f059796..e83b3278 100644 --- a/src/StakeRegistry.sol +++ b/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 diff --git a/src/libraries/BitmapUtils.sol b/src/libraries/BitmapUtils.sol index e5f789a9..58315e08 100644 --- a/src/libraries/BitmapUtils.sol +++ b/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); } /** diff --git a/test/harnesses/BitmapUtilsWrapper.sol b/test/harnesses/BitmapUtilsWrapper.sol index 3f4e6d8c..57c2073c 100644 --- a/test/harnesses/BitmapUtilsWrapper.sol +++ b/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); } } diff --git a/test/unit/BitmapUtils.t.sol b/test/unit/BitmapUtils.t.sol index 76225aa9..fedaefb5 100644 --- a/test/unit/BitmapUtils.t.sol +++ b/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"); } } } \ No newline at end of file