Add nested n flags op

0xsequence · Jan 22, 2024 · 77acc5e · 77acc5e
1 parent 9c791a7
commit 77acc5e
Show file tree

Hide file tree

Showing 3 changed files with 125 additions and 1 deletion.
diff --git a/foundry_test/modules/utils/L2CompressorEncoder.sol b/foundry_test/modules/utils/L2CompressorEncoder.sol
@@ -219,3 +219,19 @@ function encode_abi_call(bytes4 _selector, bytes32 _v1, bytes32 _v2, bytes32 _v3
 function encode_abi_call(bytes4 _selector, bytes32 _v1, bytes32 _v2, bytes32 _v3, bytes32 _v4, bytes32 _v5, bytes32 _v6) pure returns (bytes memory) {
     return abi.encodePacked(uint8(0x33), _selector, encodeWord(_v1), encodeWord(_v2), encodeWord(_v3), encodeWord(_v4), encodeWord(_v5), encodeWord(_v6));
 }
+
+function encode_nested(bytes memory _a, bytes memory _b) pure returns (bytes memory) {
+  return abi.encodePacked(uint8(0x34), uint8(2), _a, _b);
+}
+
+function encode_nested(bytes memory _a, bytes memory _b, bytes memory _c) pure returns (bytes memory) {
+  return abi.encodePacked(uint8(0x34), uint8(3), _a, _b, _c);
+}
+
+function encode_nested(bytes[] memory _bs) pure returns (bytes memory) {
+  bytes memory res = abi.encodePacked(uint8(0x35), uint16(_bs.length));
+  for (uint256 i = 0; i < _bs.length; i++) {
+    res = abi.encodePacked(res, _bs[i]);
+  }
+  return res;
+}
diff --git a/foundry_test/modules/utils/L2CompressorHuffReadFlag.t.sol b/foundry_test/modules/utils/L2CompressorHuffReadFlag.t.sol
@@ -357,4 +357,47 @@ contract L2CompressorHuffReadFlagTests is AdvTest {
     assertEq(windex, FMS + res.length);
     assertEq(abi.encodePacked(_selector, _v1, _v2, _v3, _v4, _v5, _v6), res);
   }
+
+  function test_read_nested(bytes memory _dynamic, uint256 _val1, uint256 _val2) external {
+    bytes memory encoded = encode_nested(
+      encode_bytes_n(_dynamic),
+      encode_nested(
+        encodeWord(_val1),
+        encodeWord(_val2)
+      )
+    );
+
+    (bool s, bytes memory r) = imp.staticcall(encoded);
+    assertEq(s, true);
+
+    (uint256 rindex, uint256 windex, bytes memory res) = abi.decode(r, (uint256, uint256, bytes));
+
+    assertEq(rindex, encoded.length);
+    assertEq(windex, FMS + res.length);
+    assertEq(abi.encodePacked(_dynamic, _val1, _val2), res);
+  }
+
+  function test_read_encode_nested_long() external {
+    bytes[] memory vals = new bytes[](2000);
+
+    for (uint256 i = 0; i < vals.length; i++) {
+      vals[i] = encodeWord(i * 2);
+    }
+
+    bytes memory encoded = encode_nested(vals);
+
+    (bool s, bytes memory r) = imp.staticcall(encoded);
+    assertEq(s, true);
+
+    bytes memory expected;
+    for (uint256 i = 0; i < vals.length; i++) {
+      expected = abi.encodePacked(expected, uint256(i * 2));
+    }
+
+    (uint256 rindex, uint256 windex, bytes memory res) = abi.decode(r, (uint256, uint256, bytes));
+
+    assertEq(rindex, encoded.length);
+    assertEq(windex, FMS + res.length);
+    assertEq(expected, res);
+  }
 }
diff --git a/src/L2Compressor.huff b/src/L2Compressor.huff
@@ -140,9 +140,11 @@
   FLAG_ABI_4_PARAMS          // 0x31
   FLAG_ABI_5_PARAMS          // 0x32
   FLAG_ABI_6_PARAMS          // 0x33
+  FLAG_NESTED_N_FLAGS_8      // 0x34
+  FLAG_NESTED_N_FLAGS_16     // 0x35
 }
 
-#define constant HIGHEST_FLAG = 0x33
+#define constant HIGHEST_FLAG = 0x35
 
 #define macro READ_FLAG() = takes (2) returns (2) {
   nrfs:
@@ -342,6 +344,14 @@
   FLAG_ABI_6_PARAMS:
     READ_ABI_6(nrfs) // [windex, rindex]
     end jump
+
+  FLAG_NESTED_N_FLAGS_8:
+    READ_NESTED_N_FLAGS_8(nrfs)  // [windex, rindex]
+    end jump 
+
+  FLAG_NESTED_N_FLAGS_16:
+    READ_NESTED_N_FLAGS_16(nrfs)  // [windex, rindex]
+    end jump 
 
   default:
     // The default just pushes the flag as a byte (padded to 32 bytes)
@@ -394,6 +404,61 @@
   // output stack: [mem[windex - 0x20], windex - 0x20]
 }
 
+#define macro READ_NESTED_N_FLAGS_8(nrfs) = takes (2) returns (2) {
+  // input stack: [windex, rindex]
+
+  swap1       // [rindex, windex]
+  LOAD_DYNAMIC_SIZE(0x01, 0xf8)  // [size, rindex, windex]
+
+  swap2       // [windex, rindex, size]
+  READ_NESTED_N_FLAGS(nrfs)
+
+  // output stack: [windex, rindex]
+}
+
+#define macro READ_NESTED_N_FLAGS_16(nrfs) = takes (2) returns (2) {
+  // input stack: [windex, rindex]
+
+  swap1       // [rindex, windex]
+  LOAD_DYNAMIC_SIZE(0x02, 0xf0)  // [size, rindex, windex]
+
+  swap2       // [windex, rindex, size]
+  READ_NESTED_N_FLAGS(nrfs)
+
+  // output stack: [windex, rindex]
+}
+
+#define macro READ_NESTED_N_FLAGS(nrfs) = takes (3) returns (2) {
+  // input stack: [windex, rindex, n]
+
+  0x00          // [i, windex, rindex, n]
+
+  read_more:    // [i, windex, rindex, n]
+
+    swap2       // [rindex, windex, i, n]
+    swap1       // [windex, rindex, i, n]
+
+    PERFORM_NESTED_READ_FLAG(nrfs) // [windex, rindex, i, n]
+
+    swap1                          // [rindex, windex, i, n]
+    swap2                          // [i, windex, rindex, n]
+
+    0x01 add                       // [i + 1, windex, rindex, n]
+
+    dup4                           // [n, i, windex, rindex, n]
+    dup2                           // [i, n, i, windex, rindex, n]
+    lt                             // [(i < n), i, windex, rindex, n]
+
+    read_more jumpi                // [i, windex, rindex, n]
+
+  pop           // [windex, rindex, n]
+  swap2         // [n, rindex, windex]
+  pop           // [rindex, windex]
+  swap1         // [windex, rindex]
+
+  // output stack: [windex, rindex]
+}
+
 #[calldata("0x02f1f2")]
 #define test TEST_READ_FLAG_2_BYTES() = {
   0x00  // [rindex]