diff --git a/src/libraries/BitmapUtils.sol b/src/libraries/BitmapUtils.sol index fa5e37cf..028ab60a 100644 --- a/src/libraries/BitmapUtils.sol +++ b/src/libraries/BitmapUtils.sol @@ -300,6 +300,14 @@ library BitmapUtils { return (((bitmap >> numberToCheckForInclusion) & 1) == 1); } + /** + * @notice Sets the bit at `numberToAdd` in `bitmap` to 1 and returns the result. + * Note that if the bit is already set, this function will return the same bitmap. + */ + function addNumberToBitmap(uint256 bitmap, uint8 numberToAdd) internal pure returns (uint256) { + return bitmap | (1 << numberToAdd); + } + /** * @notice Returns true if `bitmap` has no set bits */ diff --git a/test/harnesses/BitmapUtilsWrapper.sol b/test/harnesses/BitmapUtilsWrapper.sol index 77c5d917..3d2a3e3e 100644 --- a/test/harnesses/BitmapUtilsWrapper.sol +++ b/test/harnesses/BitmapUtilsWrapper.sol @@ -37,6 +37,10 @@ contract BitmapUtilsWrapper { return BitmapUtils.numberIsInBitmap(bitmap, numberToCheckForInclusion); } + function addNumberToBitmap(uint256 bitmap, uint8 numberToAdd) external pure returns (uint256) { + return BitmapUtils.addNumberToBitmap(bitmap, numberToAdd); + } + function isEmpty(uint256 bitmap) external pure returns (bool) { return BitmapUtils.isEmpty(bitmap); } diff --git a/test/unit/BitmapUtils.t.sol b/test/unit/BitmapUtils.t.sol index b0fdc61d..0c62d927 100644 --- a/test/unit/BitmapUtils.t.sol +++ b/test/unit/BitmapUtils.t.sol @@ -31,7 +31,7 @@ contract BitmapUtilsUnitTests_bitwiseOperations is BitmapUtilsUnitTests { } /// @notice some simple sanity checks on the `numberIsInBitmap` function - function test_NumberIsInBitmap() public { + function test_numberIsInBitmap() public { assertTrue(bitmapUtilsWrapper.numberIsInBitmap(2 ** 6, 6), "numberIsInBitmap function is broken 0"); assertTrue(bitmapUtilsWrapper.numberIsInBitmap(1, 0), "numberIsInBitmap function is broken 1"); assertTrue(bitmapUtilsWrapper.numberIsInBitmap(255, 7), "numberIsInBitmap function is broken 2"); @@ -42,6 +42,15 @@ contract BitmapUtilsUnitTests_bitwiseOperations is BitmapUtilsUnitTests { } } + function test_addNumberToBitmap(uint256 bitmap, uint8 numberToAdd) public { + // Ensure that numberToAdd isn't already in the bitmap + cheats.assume(bitmap | (1 << numberToAdd) != bitmap); + uint256 updatedBitmap = bitmapUtilsWrapper.addNumberToBitmap(bitmap, numberToAdd); + assertTrue( + bitmapUtilsWrapper.numberIsInBitmap(updatedBitmap, numberToAdd), "addNumberToBitmap function is broken" + ); + } + function testFuzz_isEmpty(uint256 input) public { if (input == 0) { // assertTrue(bitmapUtilsWrapper.isEmpty(input), "isEmpty function is broken");