Bug 1854007 - wasm: Fix serialization of duplicate recursion groups with self references. r=bvisness

We store a *typedef -> u32 map on TypeContext which we can use to get the type index
for a type definition. When a module has multiple types that canonicalize to the
same value, we only store the first index in the map.

We use this map when serializing the type section. We currently serialize every
recursion group (even if it was already serialized). When deserializing we re-run
canonicalization which will de-duplicate this.

These two combined can cause a module with types:
 0: (type (struct (field 0)))
 1: (type (struct (field 1))) ;; identical to 0

To be serialized as:
 0: (type (struct (field 0)))
 1: (type (struct (field 0))) ;; not identical to 0!

The solution is to only encode the first recursion group and then
re-use it for the second recursion group.

Differential Revision: https://phabricator.services.mozilla.com/D188644

js/src/jit-test/tests/wasm/gc/bug-1854007.js

@@ -0,0 +1,16 @@
+// |jit-test| test-also=--wasm-test-serialization; skip-if: !wasmGcEnabled()
+
+let {run} = wasmEvalText(`(module
+  (rec (type $$t1 (func (result (ref null $$t1)))))
+  (rec (type $$t2 (func (result (ref null $$t2)))))
+
+  (func $$f1 (type $$t1) (ref.null $$t1))
+  (func $$f2 (type $$t2) (ref.null $$t2))
+  (table funcref (elem $$f1 $$f2))
+
+  (func (export "run")
+    (call_indirect (type $$t2) (i32.const 0))
+    drop
+  )
+)`).exports;
+run();

js/src/wasm/WasmSerialize.cpp

@@ -598,6 +598,10 @@ CoderResult CodeTypeDef(Coder<mode>& coder, CoderArg<mode, TypeDef> item) {
   return Ok();
 }
 
+using RecGroupIndexMap =
+    HashMap<const RecGroup*, uint32_t, PointerHasher<const RecGroup*>,
+            SystemAllocPolicy>;
+
 template <CoderMode mode>
 CoderResult CodeTypeContext(Coder<mode>& coder,
                             CoderArg<mode, TypeContext> item) {
@@ -614,6 +618,22 @@ CoderResult CodeTypeContext(Coder<mode>& coder,
     // Decode each recursion group
     for (uint32_t recGroupIndex = 0; recGroupIndex < numRecGroups;
          recGroupIndex++) {
+      // Decode if this recursion group is equivalent to a previous recursion
+      // group
+      uint32_t canonRecGroupIndex;
+      MOZ_TRY(CodePod(coder, &canonRecGroupIndex));
+      MOZ_RELEASE_ASSERT(canonRecGroupIndex <= recGroupIndex);
+
+      // If the decoded index is not ours, we must re-use the previous decoded
+      // recursion group.
+      if (canonRecGroupIndex != recGroupIndex) {
+        SharedRecGroup recGroup = item->groups()[canonRecGroupIndex];
+        if (!item->addRecGroup(recGroup)) {
+          return Err(OutOfMemory());
+        }
+        continue;
+      }
+
       // Decode the number of types in the recursion group
       uint32_t numTypes;
       MOZ_TRY(CodePod(coder, &numTypes));
@@ -639,10 +659,45 @@ CoderResult CodeTypeContext(Coder<mode>& coder,
     uint32_t numRecGroups = item->groups().length();
     MOZ_TRY(CodePod(coder, &numRecGroups));
 
+    // We must be careful to only encode every unique recursion group only once
+    // and in module order. The reason for this is that encoding type def
+    // references uses the module type index map, which only stores the first
+    // type index a type was canonicalized to.
+    //
+    // Using this map to encode both recursion groups would turn the following
+    // type section from:
+    //
+    // 0: (type (struct (field 0)))
+    // 1: (type (struct (field 1))) ;; identical to 0
+    //
+    // into:
+    //
+    // 0: (type (struct (field 0)))
+    // 1: (type (struct (field 0))) ;; not identical to 0!
+    RecGroupIndexMap canonRecGroups;
+
     // Encode each recursion group
     for (uint32_t groupIndex = 0; groupIndex < numRecGroups; groupIndex++) {
       SharedRecGroup group = item->groups()[groupIndex];
 
+      // Find the index of the first time this recursion group was encoded, or
+      // set it to this index if it hasn't been encoded.
+      RecGroupIndexMap::AddPtr canonRecGroupIndex =
+          canonRecGroups.lookupForAdd(group.get());
+      if (!canonRecGroupIndex) {
+        if (!canonRecGroups.add(canonRecGroupIndex, group.get(), groupIndex)) {
+          return Err(OutOfMemory());
+        }
+      }
+
+      // Encode the canon index for this recursion group
+      MOZ_TRY(CodePod(coder, &canonRecGroupIndex->value()));
+
+      // Don't encode this recursion group if we've already encoded it
+      if (canonRecGroupIndex->value() != groupIndex) {
+        continue;
+      }
+
       // Encode the number of types in the recursion group
       uint32_t numTypes = group->numTypes();
       MOZ_TRY(CodePod(coder, &numTypes));

js/src/wasm/WasmTypeDef.h

@@ -1137,24 +1137,14 @@ class TypeContext : public AtomicRefCounted<TypeContext> {
     MOZ_ASSERT(!pendingRecGroup_);
 
     // Create the group and add it to the list of groups
-    pendingRecGroup_ = RecGroup::allocate(numTypes);
-    if (!pendingRecGroup_ || !recGroups_.append(pendingRecGroup_)) {
+    MutableRecGroup recGroup = RecGroup::allocate(numTypes);
+    if (!recGroup || !addRecGroup(recGroup)) {
       return nullptr;
     }
 
-    // Store the types of the group into our index space maps. These may get
-    // overwritten when we finish this group and canonicalize it. We need to do
-    // this before finishing, because these entries will be used by decoding
-    // and error printing.
-    for (uint32_t groupTypeIndex = 0; groupTypeIndex < numTypes;
-         groupTypeIndex++) {
-      const TypeDef* typeDef = &pendingRecGroup_->type(groupTypeIndex);
-      uint32_t typeIndex = types_.length();
-      if (!types_.append(typeDef) || !moduleIndices_.put(typeDef, typeIndex)) {
-        return nullptr;
-      }
-    }
-    return pendingRecGroup_;
+    // Store this group for later use in endRecGroup
+    pendingRecGroup_ = recGroup;
+    return recGroup;
   }
 
   // Finish creation of a recursion group after type definitions have been
@@ -1211,6 +1201,33 @@ class TypeContext : public AtomicRefCounted<TypeContext> {
     return true;
   }
 
+  // Finish creation of a recursion group after type definitions have been
+  // initialized. This must be paired with `startGroup`.
+  [[nodiscard]] bool addRecGroup(SharedRecGroup recGroup) {
+    // We must not have a pending group
+    MOZ_ASSERT(!pendingRecGroup_);
+
+    // Add it to the list of groups
+    if (!recGroups_.append(recGroup)) {
+      return false;
+    }
+
+    // Store the types of the group into our index space maps. These may get
+    // overwritten if this group is being added by `startRecGroup` and we
+    // overwrite it with a canonical group in `endRecGroup`. We need to do
+    // this before finishing though, because these entries will be used by
+    // decoding and error printing.
+    for (uint32_t groupTypeIndex = 0; groupTypeIndex < recGroup->numTypes();
+         groupTypeIndex++) {
+      const TypeDef* typeDef = &recGroup->type(groupTypeIndex);
+      uint32_t typeIndex = types_.length();
+      if (!types_.append(typeDef) || !moduleIndices_.put(typeDef, typeIndex)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   template <typename T>
   [[nodiscard]] bool addType(T&& type) {
     MutableRecGroup recGroup = startRecGroup(1);