Isorecursive binary format (#4494)

Write and parse recursion groups in binary type sections. Unlike in the text
format, where we ignore recursion groups when not using isorecursive types, do
not allow parsing binary recursion group when using other type systems. Doing so
would produce incorrect results because recursions groups only count as single
entries in the type system vector so we dynamically grow the TypeBuilder when we
encounter them. That would change the mapping of later indices to types, and
would change the meaning of previous type definitions that use those later
indices. This is not a problem in the isorecursive system because in that system
type definitions are not allowed to use later indices.

6 files changed, 57 insertions, 18 deletions

diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index 0157f4858..b31447ed9 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -386,6 +386,8 @@ enum EncodedType {
   FuncExtending = -0x23,   // 0x5d
   StructExtending = -0x24, // 0x5c
   ArrayExtending = -0x25,  // 0x5b
+  // isorecursive recursion groups
+  Rec = -0x31, // 0x4f
   // block_type
   Empty = -0x40 // 0x40
 };
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 13bcc5110..1ddb009cf 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -219,16 +219,38 @@ void WasmBinaryWriter::writeTypes() {
   if (indexedTypes.types.size() == 0) {
     return;
   }
+  // Count the number of recursion groups, which is always the number of
+  // elements in the type section. In non-isorecursive type systems, it is also
+  // equivalent to the number of types.
+  std::optional<RecGroup> lastGroup;
+  size_t numGroups = 0;
+  for (auto type : indexedTypes.types) {
+    auto currGroup = type.getRecGroup();
+    numGroups += lastGroup != currGroup;
+    lastGroup = currGroup;
+  }
   BYN_TRACE("== writeTypes\n");
   auto start = startSection(BinaryConsts::Section::Type);
-  o << U32LEB(indexedTypes.types.size());
+  o << U32LEB(numGroups);
+  lastGroup = std::nullopt;
   for (Index i = 0; i < indexedTypes.types.size(); ++i) {
     auto type = indexedTypes.types[i];
-    bool nominal = getTypeSystem() == TypeSystem::Nominal;
+    // Check whether we need to start a new recursion group. Recursion groups of
+    // size 1 are implicit, so only emit a group header for larger groups. This
+    // gracefully handles non-isorecursive type systems, which only have groups
+    // of size 1.
+    auto currGroup = type.getRecGroup();
+    if (lastGroup != currGroup && currGroup.size() > 1) {
+      o << S32LEB(BinaryConsts::EncodedType::Rec) << U32LEB(currGroup.size());
+    }
+    lastGroup = currGroup;
+    // Emit the type definition.
+    bool hasSupertype = getTypeSystem() == TypeSystem::Nominal ||
+                        getTypeSystem() == TypeSystem::Isorecursive;
     BYN_TRACE("write " << type << std::endl);
     if (type.isSignature()) {
-      o << S32LEB(nominal ? BinaryConsts::EncodedType::FuncExtending
-                          : BinaryConsts::EncodedType::Func);
+      o << S32LEB(hasSupertype ? BinaryConsts::EncodedType::FuncExtending
+                               : BinaryConsts::EncodedType::Func);
       auto sig = type.getSignature();
       for (auto& sigType : {sig.params, sig.results}) {
         o << U32LEB(sigType.size());
@@ -237,21 +259,21 @@ void WasmBinaryWriter::writeTypes() {
         }
       }
     } else if (type.isStruct()) {
-      o << S32LEB(nominal ? BinaryConsts::EncodedType::StructExtending
-                          : BinaryConsts::EncodedType::Struct);
+      o << S32LEB(hasSupertype ? BinaryConsts::EncodedType::StructExtending
+                               : BinaryConsts::EncodedType::Struct);
       auto fields = type.getStruct().fields;
       o << U32LEB(fields.size());
       for (const auto& field : fields) {
         writeField(field);
       }
     } else if (type.isArray()) {
-      o << S32LEB(nominal ? BinaryConsts::EncodedType::ArrayExtending
-                          : BinaryConsts::EncodedType::Array);
+      o << S32LEB(hasSupertype ? BinaryConsts::EncodedType::ArrayExtending
+                               : BinaryConsts::EncodedType::Array);
       writeField(type.getArray().element);
     } else {
       WASM_UNREACHABLE("TODO GC type writing");
     }
-    if (nominal) {
+    if (hasSupertype) {
       auto super = type.getSuperType();
       if (!super) {
         super = type.isFunction() ? HeapType::func : HeapType::data;
@@ -1844,9 +1866,8 @@ void WasmBinaryBuilder::readMemory() {
 
 void WasmBinaryBuilder::readTypes() {
   BYN_TRACE("== readTypes\n");
-  size_t numTypes = getU32LEB();
-  BYN_TRACE("num: " << numTypes << std::endl);
-  TypeBuilder builder(numTypes);
+  TypeBuilder builder(getU32LEB());
+  BYN_TRACE("num: " << builder.size() << std::endl);
 
   auto makeType = [&](int32_t typeCode) {
     Type type;
@@ -1865,7 +1886,7 @@ void WasmBinaryBuilder::readTypes() {
         if (getBasicHeapType(htCode, ht)) {
           return Type(ht, nullability);
         }
-        if (size_t(htCode) >= numTypes) {
+        if (size_t(htCode) >= builder.size()) {
           throwError("invalid type index: " + std::to_string(htCode));
         }
         return builder.getTempRefType(builder[size_t(htCode)], nullability);
@@ -1875,7 +1896,7 @@ void WasmBinaryBuilder::readTypes() {
         auto depth = typeCode == BinaryConsts::EncodedType::rtt ? Rtt::NoDepth
                                                                 : getU32LEB();
         auto htCode = getU32LEB();
-        if (size_t(htCode) >= numTypes) {
+        if (size_t(htCode) >= builder.size()) {
           throwError("invalid type index: " + std::to_string(htCode));
         }
         return builder.getTempRttType(Rtt(depth, builder[htCode]));
@@ -1944,9 +1965,23 @@ void WasmBinaryBuilder::readTypes() {
     return Struct(std::move(fields));
   };
 
-  for (size_t i = 0; i < numTypes; i++) {
+  for (size_t i = 0; i < builder.size(); i++) {
     BYN_TRACE("read one\n");
     auto form = getS32LEB();
+    if (form == BinaryConsts::EncodedType::Rec) {
+      if (getTypeSystem() != TypeSystem::Isorecursive) {
+        Fatal() << "Binary recursion groups only supported in --hybrid mode";
+      }
+      uint32_t groupSize = getU32LEB();
+      if (groupSize == 0u) {
+        Fatal() << "Invalid recursion group of size zero";
+      }
+      // The group counts as one element in the type section, so we have to
+      // allocate space for the extra types.
+      builder.grow(groupSize - 1);
+      builder.createRecGroup(i, groupSize);
+      form = getS32LEB();
+    }
     if (form == BinaryConsts::EncodedType::Func ||
         form == BinaryConsts::EncodedType::FuncExtending) {
       builder[i] = readSignatureDef();
@@ -1964,7 +1999,7 @@ void WasmBinaryBuilder::readTypes() {
         form == BinaryConsts::EncodedType::ArrayExtending) {
       auto superIndex = getS64LEB(); // TODO: Actually s33
       if (superIndex >= 0) {
-        if (size_t(superIndex) >= numTypes) {
+        if (size_t(superIndex) >= builder.size()) {
           throwError("bad supertype index " + std::to_string(superIndex));
         }
         builder[i].subTypeOf(builder[superIndex]);
diff --git a/test/lit/isorecursive-good.wast b/test/lit/isorecursive-good.wast
index 353905c4f..7ced0d8d3 100644
--- a/test/lit/isorecursive-good.wast
+++ b/test/lit/isorecursive-good.wast
@@ -1,10 +1,9 @@
 ;; TODO: Autogenerate these checks! The current script cannot handle `rec`.
 
 ;; RUN: wasm-opt %s -all --hybrid -S -o - | filecheck %s --check-prefix HYBRID
+;; RUN: wasm-opt %s -all --hybrid --roundtrip -S -o - | filecheck %s --check-prefix HYBRID
 ;; RUN: wasm-opt %s -all --nominal -S -o - | filecheck %s --check-prefix NOMINAL
 
-;; TDOO: test with --roundtrip as well
-
 (module
 
 ;; HYBRID:      (rec
diff --git a/test/lit/isorecursive-output-ordering.wast b/test/lit/isorecursive-output-ordering.wast
index 6aaa7ba39..ae965123f 100644
--- a/test/lit/isorecursive-output-ordering.wast
+++ b/test/lit/isorecursive-output-ordering.wast
@@ -1,6 +1,7 @@
 ;; TODO: Autogenerate these checks! The current script cannot handle `rec`.
 
 ;; RUN: foreach %s %t wasm-opt -all --hybrid -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt -all --hybrid --roundtrip -S -o - | filecheck %s
 
 (module
  ;; Test that we order groups by average uses.
diff --git a/test/lit/isorecursive-singleton-group.wast b/test/lit/isorecursive-singleton-group.wast
index adb40b141..eeb92ac09 100644
--- a/test/lit/isorecursive-singleton-group.wast
+++ b/test/lit/isorecursive-singleton-group.wast
@@ -1,6 +1,7 @@
 ;; TODO: Autogenerate these checks! The current script cannot handle `rec`.
 
 ;; RUN: wasm-opt %s -all --hybrid -S -o - | filecheck %s
+;; RUN: wasm-opt %s -all --hybrid --roundtrip -S -o - | filecheck %s
 
 ;; Check that everything works correctly when a recursion group has only a
 ;; single member. The rec group is implicit, so does not need to be printed.
diff --git a/test/lit/isorecursive-whole-group.wast b/test/lit/isorecursive-whole-group.wast
index 625025b23..4d349a3cb 100644
--- a/test/lit/isorecursive-whole-group.wast
+++ b/test/lit/isorecursive-whole-group.wast
@@ -1,6 +1,7 @@
 ;; TODO: Autogenerate these checks! The current script cannot handle `rec`.
 
 ;; RUN: wasm-opt %s -all --hybrid -S -o - | filecheck %s
+;; RUN: wasm-opt %s -all --hybrid --roundtrip -S -o - | filecheck %s
 
 ;; Check that unused types are still included in the output when they are part
 ;; of a recursion group with used types.