feat: addNumberToBitmap function

Layr-Labs · Dec 12, 2023 · 9e4fc83 · 9e4fc83
1 parent 6dbe88b
commit 9e4fc83
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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");