Eagerly canonicalize basic types (#4507)

We were already eagerly canonicalizing basic HeapTypes when building types so
the more complicated canonicalization algorithms would not have to handle
noncanonical heap types, but we were not doing the same for Types. Equirecursive
canonicalization was properly handling noncanonical Types everywhere, but
isorecursive canonicalization was not. Rather than update isorecursive
canonicalization in multiple places, remove the special handling from
equirecursive canonicalization and canonicalize types in a single location.

2 files changed, 48 insertions, 8 deletions

diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 04b1122b6..d7ba224c8 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -2194,7 +2194,6 @@ std::ostream& TypePrinter::print(const Rtt& rtt) {
 }
 
 size_t FiniteShapeHasher::hash(Type type) {
-  type = asCanonical(type);
   size_t digest = wasm::hash(type.isBasic());
   if (type.isBasic()) {
     rehash(digest, type.getID());
@@ -2205,7 +2204,6 @@ size_t FiniteShapeHasher::hash(Type type) {
 }
 
 size_t FiniteShapeHasher::hash(HeapType heapType) {
-  heapType = asCanonical(heapType);
   size_t digest = wasm::hash(heapType.isBasic());
   if (heapType.isBasic()) {
     rehash(digest, heapType.getID());
@@ -2314,8 +2312,6 @@ size_t FiniteShapeHasher::hash(const Rtt& rtt) {
 }
 
 bool FiniteShapeEquator::eq(Type a, Type b) {
-  a = asCanonical(a);
-  b = asCanonical(b);
   if (a.isBasic() != b.isBasic()) {
     return false;
   } else if (a.isBasic()) {
@@ -2326,8 +2322,6 @@ bool FiniteShapeEquator::eq(Type a, Type b) {
 }
 
 bool FiniteShapeEquator::eq(HeapType a, HeapType b) {
-  a = asCanonical(a);
-  b = asCanonical(b);
   if (a.isBasic() != b.isBasic()) {
     return false;
   } else if (a.isBasic()) {
@@ -3798,7 +3792,7 @@ std::optional<TypeBuilder::Error> canonicalizeIsorecursive(
   return {};
 }
 
-void canonicalizeBasicHeapTypes(CanonicalizationState& state) {
+void canonicalizeBasicTypes(CanonicalizationState& state) {
   // Replace heap types backed by BasicKind HeapTypeInfos with their
   // corresponding BasicHeapTypes. The heap types backed by BasicKind
   // HeapTypeInfos exist only to support building basic types in a TypeBuilder
@@ -3814,6 +3808,25 @@ void canonicalizeBasicHeapTypes(CanonicalizationState& state) {
     }
   }
   state.update(replacements);
+
+  if (replacements.size()) {
+    // Canonicalizing basic heap types may cause their parent types to become
+    // canonicalizable as well, for example after creating `(ref null extern)`
+    // we can futher canonicalize to `externref`.
+    struct TypeCanonicalizer : TypeGraphWalkerBase<TypeCanonicalizer> {
+      void scanType(Type* type) {
+        if (type->isTuple()) {
+          TypeGraphWalkerBase<TypeCanonicalizer>::scanType(type);
+        } else {
+          *type = asCanonical(*type);
+        }
+      }
+    };
+    for (auto& info : state.newInfos) {
+      auto root = asHeapType(info);
+      TypeCanonicalizer{}.walkRoot(&root);
+    }
+  }
 }
 
 } // anonymous namespace
@@ -3842,7 +3855,7 @@ TypeBuilder::BuildResult TypeBuilder::build() {
 
   // Eagerly replace references to built basic heap types so the more
   // complicated canonicalization algorithms don't need to consider them.
-  canonicalizeBasicHeapTypes(state);
+  canonicalizeBasicTypes(state);
 
 #if TRACE_CANONICALIZATION
   std::cerr << "After replacing basic heap types:\n";
diff --git a/test/gtest/type-builder.cpp b/test/gtest/type-builder.cpp
index 2a29d11b0..4ae9e0e63 100644
--- a/test/gtest/type-builder.cpp
+++ b/test/gtest/type-builder.cpp
@@ -457,3 +457,30 @@ TEST_F(IsorecursiveTest, CanonicalizeTypesBeforeSubtyping) {
   auto result = builder.build();
   EXPECT_TRUE(result);
 }
+
+static void testCanonicalizeBasicTypes() {
+  TypeBuilder builder(5);
+
+  Type externref = builder.getTempRefType(builder[0], Nullable);
+  Type externrefs = builder.getTempTupleType({externref, externref});
+
+  builder[0] = HeapType::ext;
+  builder[1] = Struct({Field(externref, Immutable)});
+  builder[2] = Struct({Field(Type::externref, Immutable)});
+  builder[3] = Signature(externrefs, Type::none);
+  builder[4] = Signature({Type::externref, Type::externref}, Type::none);
+
+  auto result = builder.build();
+  ASSERT_TRUE(result);
+  auto built = *result;
+
+  EXPECT_EQ(built[1], built[2]);
+  EXPECT_EQ(built[3], built[4]);
+}
+
+TEST_F(EquirecursiveTest, CanonicalizeBasicTypes) {
+  testCanonicalizeBasicTypes();
+}
+TEST_F(IsorecursiveTest, CanonicalizeBasicTypes) {
+  testCanonicalizeBasicTypes();
+}