cleanup

nfurfaro · Oct 14, 2023 · c29216b · c29216b
1 parent 01a2329
commit c29216b
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 34 deletions.
diff --git a/src/Huffbits.huff b/src/Huffbits.huff
@@ -6,18 +6,19 @@
 
 
 // Interface
-#define function singleBitMask(uint256) pure returns(uint256)
+#define function bitMask(uint256) pure returns(uint256)
 #define function multiMask(uint256) pure returns(uint256)
 #define function nibbleMask(uint256) pure returns(uint256)
 #define function toggleBit(uint256, uint256) pure returns(uint256)
 #define function queryBit(uint256, uint256) pure returns(uint256)
 #define function setBit(uint256, uint256) pure returns(uint256)
 #define function clearBit(uint256, uint256) pure returns(uint256)
+#define function clearNibble(uint256, uint256) pure returns(uint256)
 
 #define constant MAX = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF
 
 /// Get a bitmask with a single `1` at the nth position.
-#define macro SINGLE_BITMASK() = takes(1) returns(1) {
+#define macro BITMASK() = takes(1) returns(1) {
     // input stack          //                 [shift]
     0x01                    //              [1, shift]
     swap1                   //              [shift, 1]
@@ -80,7 +81,7 @@ pub fn b256_multimask(n: u64) -> b256 {
 /// Flip the nth bit in a bitmap.
 #define macro TOGGLE_BIT() = takes(2) returns(1) {
     // input stack          //       [bit, bitmap]
-    SINGLE_BITMASK()        //      [mask, bitmap]
+    BITMASK()        //      [mask, bitmap]
     xor                     //            [result]
 }
 
@@ -90,7 +91,7 @@ pub fn b256_multimask(n: u64) -> b256 {
     dup1                   //  [bit, bit, bitmap]
     swap2                  //  [bitmap, bit, bit]
     swap1                  //  [bit, bitmap, bit]
-    SINGLE_BITMASK()       // [mask, bitmap, bit]
+    BITMASK()       // [mask, bitmap, bit]
     and                    //          [and, bit]
     swap1                  //          [bit, and]
     shr                    //           [shifted]
@@ -112,4 +113,11 @@ pub fn b256_multimask(n: u64) -> b256 {
     swap1              // [bit, full_map, bitmap]
     TOGGLE_BIT()       //          [mask, bitmap]
     and
+}
+
+/// Given a bitmap, clear a nibble if set, otherwise return the original bitmap.
+#define macro CLEAR_NIBBLE() = takes(2) returns(1) {
+    // input stack     //            [index, bitmap]
+    NIBBLEMASK()       //             [mask, bitmap]
+    and
 }
diff --git a/test/Huffbits.t.sol b/test/Huffbits.t.sol
@@ -6,13 +6,14 @@ import "forge-std/Test.sol";
 import "forge-std/console.sol";
 
 interface Huffbits {
-    function singleBitMask(uint256) external pure returns(uint256);
+    function bitMask(uint256) external pure returns(uint256);
     function multiMask(uint256) external pure returns(uint256);
     function nibbleMask(uint256) external pure returns(uint256);
     function toggleBit(uint256, uint256) external pure returns(uint256);
     function queryBit(uint256, uint256) external pure returns(uint256);
     function setBit(uint256, uint256) external pure returns(uint256);
     function clearBit(uint256, uint256) external pure returns(uint256);
+    function clearNibble(uint256, uint256) external pure returns(uint256);
 }
 
 
@@ -27,15 +28,15 @@ contract HuffbitsTest is Test {
     }
 
     /// @dev Ensure that a single bit can be set.
-    function testSingleBitMask() public {
-        assertEq(huffbits.singleBitMask(0x0), 0x1);
-        assertEq(huffbits.singleBitMask(0x1), 0x2);
-        assertEq(huffbits.singleBitMask(0x2), 0x4);
-        assertEq(huffbits.singleBitMask(0x3), 0x8);
-        assertEq(huffbits.singleBitMask(0x8), 0x100);
-        assertEq(huffbits.singleBitMask(0x8), 0x100);
-        assertEq(huffbits.singleBitMask(0x20), 0x100000000);
-        assertEq(huffbits.singleBitMask(0x3F), 0x8000000000000000);
+    function testBitMask() public {
+        assertEq(huffbits.bitMask(0x0), 0x1);
+        assertEq(huffbits.bitMask(0x1), 0x2);
+        assertEq(huffbits.bitMask(0x2), 0x4);
+        assertEq(huffbits.bitMask(0x3), 0x8);
+        assertEq(huffbits.bitMask(0x8), 0x100);
+        assertEq(huffbits.bitMask(0x8), 0x100);
+        assertEq(huffbits.bitMask(0x20), 0x100000000);
+        assertEq(huffbits.bitMask(0x3F), 0x8000000000000000);
     }
 
     function testToggleBit() public {
@@ -112,5 +113,15 @@ contract HuffbitsTest is Test {
         assertEq(huffbits.nibbleMask(55), 0xFFFFFFFF0FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF);
         assertEq(huffbits.nibbleMask(63), 0x0FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF);
     }
+
+    function testClearNibble() public {
+        assertEq(huffbits.clearNibble(0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF, 0), 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF0);
+        assertEq(huffbits.clearNibble(1, 0), 0);
+
+    }
+
+    // writeNibble
+    // queryNibble
+
 }
 
diff --git a/test/mocks/HuffbitsWrappers.huff b/test/mocks/HuffbitsWrappers.huff
@@ -1,7 +1,7 @@
 
-#define macro SINGLE_BITMASK_WRAPPER() = {
+#define macro BITMASK_WRAPPER() = {
     0x04 calldataload
-    SINGLE_BITMASK()
+    BITMASK()
     0x00 mstore
     0x20 0x00 return
 }
@@ -29,44 +29,53 @@
 }
 
 #define macro QUERY_BIT_WRAPPER() = {
-    0x04 calldataload  //   [1]
-    0x24 calldataload  // [0,1]
-    QUERY_BIT()        // [set?]
-    0x00 mstore        // [],<set?>
-    0x20 0x00 return   // [],<>
+    0x04 calldataload  //      [arg1]
+    0x24 calldataload  //[arg2, arg1]
+    QUERY_BIT()        //      [set?]
+    0x00 mstore
+    0x20 0x00 return
 }
 
 #define macro SET_BIT_WRAPPER() = {
-    0x04 calldataload  //   [1]
-    0x24 calldataload  // [0,1]
-    SET_BIT()        // [set?]
-    0x00 mstore        // [],<set?>
-    0x20 0x00 return   // [],<>
+    0x04 calldataload//      [arg1]
+    0x24 calldataload//[arg2, arg1]
+    SET_BIT()        //       [set]
+    0x00 mstore
+    0x20 0x00 return
 }
 
 #define macro CLEAR_BIT_WRAPPER() = {
-    0x04 calldataload  //   [1]
-    0x24 calldataload  // [0,1]
-    CLEAR_BIT()        // [set?]
-    0x00 mstore        // [],<set?>
-    0x20 0x00 return   // [],<>
+    0x04 calldataload  //      [arg1]
+    0x24 calldataload  //[arg2, arg1]
+    CLEAR_BIT()        //   [cleared]
+    0x00 mstore
+    0x20 0x00 return
+}
+
+#define macro CLEAR_NIBBLE_WRAPPER() = {
+    0x04 calldataload  //         [arg1]
+    0x24 calldataload  //   [arg2, arg1]
+    CLEAR_NIBBLE()        //   [cleared]
+    0x00 mstore
+    0x20 0x00 return
 }
 
 #define macro MAIN() = {
     0x00 calldataload 0xE0 shr
-    dup1 __FUNC_SIG(singleBitMask) eq singleBitMask jumpi
+    dup1 __FUNC_SIG(bitMask) eq bitMask jumpi
     dup1 __FUNC_SIG(multiMask) eq multiMask jumpi
     dup1 __FUNC_SIG(nibbleMask) eq nibbleMask jumpi
     dup1 __FUNC_SIG(toggleBit) eq toggleBit jumpi
     dup1 __FUNC_SIG(queryBit) eq queryBit jumpi
     dup1 __FUNC_SIG(setBit) eq setBit jumpi
     dup1 __FUNC_SIG(clearBit) eq clearBit jumpi
+    dup1 __FUNC_SIG(clearNibble) eq clearNibble jumpi
 
 
     0x00 0x00 revert
 
-    singleBitMask:
-        SINGLE_BITMASK_WRAPPER()
+    bitMask:
+        BITMASK_WRAPPER()
     multiMask:
         MULTIMASK_WRAPPER()
     nibbleMask:
@@ -79,5 +88,7 @@
         SET_BIT_WRAPPER()
     clearBit:
         CLEAR_BIT_WRAPPER()
+    clearNibble:
+        CLEAR_NIBBLE_WRAPPER()
 
 }