Consider self-referential HeapTypes to be canonical (#3626)

This PR makes the TypeBuilder move self-referential HeapTypes to global HeapType
store so that they are considered canonical. This means that when a HeapType is
constructed with an identical definition, it will be equivalent to the original
HeapType constructed by the TypeBuilder, but it does _not_ mean that
self-referential HeapTypes are deduplicated.

This fixes a bug in which two versions of each self-referential function
signature were emitted. Before this PR, the global HeapType store was not aware
of the original self-referential HeapTypes. When the function signatures were
used to construct HeapTypes during type collection, new HeapTypes were allocated
and emitted in addition to the original HeapTypes. Now the global HeapType store
returns the original HeapTypes, so the extra HeapType is never allocated.

2 files changed, 13 insertions, 21 deletions

diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 0941f8c58..8546e3fc9 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -306,18 +306,21 @@ template<typename Info> struct Store {
   // Maps from constructed types to their canonical Type IDs.
   std::unordered_map<Info, uintptr_t> typeIDs;
 
+  TypeID recordCanonical(std::unique_ptr<Info>&& info) {
+    TypeID id = uintptr_t(info.get());
+    assert(id > Info::type_t::_last_basic_type);
+    typeIDs[*info] = id;
+    constructedTypes.emplace_back(std::move(info));
+    return id;
+  }
+
   typename Info::type_t canonicalize(const Info& info) {
     std::lock_guard<std::mutex> lock(mutex);
     auto indexIt = typeIDs.find(info);
     if (indexIt != typeIDs.end()) {
       return typename Info::type_t(indexIt->second);
     }
-    auto ptr = std::make_unique<Info>(info);
-    auto id = uintptr_t(ptr.get());
-    constructedTypes.push_back(std::move(ptr));
-    assert(id > Info::type_t::_last_basic_type);
-    typeIDs[info] = id;
-    return typename Info::type_t(id);
+    return typename Info::type_t(recordCanonical(std::make_unique<Info>(info)));
   }
 };
 
@@ -1360,10 +1363,6 @@ struct Canonicalizer {
   void canonicalize(T* type, std::unordered_map<T, T>& canonicals);
 };
 
-// Used to extend the lifetime of self-referential HeapTypes so they don't need
-// to be canonicalized.
-std::vector<std::unique_ptr<HeapTypeInfo>> noncanonicalHeapTypes;
-
 // Traverse the type graph rooted at the initialized HeapTypeInfos in reverse
 // postorder, replacing in place all Types and HeapTypes backed by the
 // TypeBuilder's Stores with equivalent globally canonicalized Types and
@@ -1413,11 +1412,12 @@ Canonicalizer::Canonicalizer(TypeBuilder& builder) : builder(builder) {
   }
 
   // Now that all other Types and HeapTypes have been canonicalized, move
-  // self-referential HeapTypes to the global store so that they will outlive
-  // the TypeBuilder without their IDs changing.
+  // self-referential HeapTypes to the global store so that they will be
+  // considered canonical and outlive the TypeBuilder.
+  std::lock_guard<std::mutex> lock(globalHeapTypeStore.mutex);
   for (auto& entry : builder.impl->entries) {
     if (selfReferentialHeapTypes.count(entry.get())) {
-      noncanonicalHeapTypes.emplace_back(std::move(entry.info));
+      globalHeapTypeStore.recordCanonical(std::move(entry.info));
     }
   }
 }
diff --git a/test/lit/recursive-types.wast b/test/lit/recursive-types.wast
index c825c9cad..7f719a359 100644
--- a/test/lit/recursive-types.wast
+++ b/test/lit/recursive-types.wast
@@ -2,20 +2,12 @@
 
 ;; RUN: wasm-opt %s -all -S -o - | filecheck %s
 
-;; TODO: Fix the bug where self-referential function signatures are emitted
-;; twice. This happens because collectHeapTypes has to reconstruct a HeapType
-;; from each function's signature, but self-referential HeapTypes aren't
-;; canonicalized when they are parsed so collectHeapTypes ends up with a
-;; separate, unfolded version of the type.
-
 ;; TODO: Fix the bug where structurally identical types are given the same
 ;; generated name, making the wast invalid due to duplicate names.
 
 ;; CHECK:      (module
 ;; CHECK-NEXT:   (type $ref?|...0|_=>_ref?|...0| (func (param (ref null $ref?|...0|_=>_ref?|...0|)) (result (ref null $ref?|...0|_=>_ref?|...0|))))
 ;; CHECK-NEXT:   (type $ref?|...0|_=>_ref?|...0| (func (param (ref null $ref?|...0|_=>_ref?|...0|)) (result (ref null $ref?|...0|_=>_ref?|...0|))))
-;; CHECK-NEXT:   (type $ref?|ref?|...0|_->_ref?|...0||_=>_ref?|ref?|...0|_->_ref?|...0|| (func (param (ref null $ref?|...0|_=>_ref?|...0|)) (result (ref null $ref?|...0|_=>_ref?|...0|))))
-;; CHECK-NEXT:   (type $ref?|ref?|...0|_->_ref?|...0||_=>_ref?|ref?|...0|_->_ref?|...0|| (func (param (ref null $ref?|...0|_=>_ref?|...0|)) (result (ref null $ref?|...0|_=>_ref?|...0|))))
 ;; CHECK-NEXT:   (func $foo (param $0 (ref null $ref?|...0|_=>_ref?|...0|)) (result (ref null $ref?|...0|_=>_ref?|...0|))
 ;; CHECK-NEXT:     (unreachable)
 ;; CHECK-NEXT:   )