Allow building basic HeapTypes in nominal mode (#4346)

As we work toward allowing nominal and structural types to coexist, any
difference in how they can be built or used will be an inconvenient footgun that
we will have to work around. In the spirit of reducing the differences between
the type systems, allow TypeBuilder to construct basic HeapTypes in nominal mode
just as it can in equirecursive mode.

Although this change is a net increase in code complexity for not much
benefit (wasm-opt never needs to build basic HeapTypes), it is also an
incremental step toward getting rid of separate type system modes, so I expect
it to simplify other PRs in the near future.

This change also uncovered a bug in how the type fuzzer generated subtypes of
basic HeapTypes. The generated subtypes did not necessarily have the intended
`Kind`, which caused failures in nominal subtype validation in the fuzzer.

2 files changed, 157 insertions, 94 deletions

diff --git a/src/tools/fuzzing/heap-types.cpp b/src/tools/fuzzing/heap-types.cpp
index 1878f5473..9992baf55 100644
--- a/src/tools/fuzzing/heap-types.cpp
+++ b/src/tools/fuzzing/heap-types.cpp
@@ -365,13 +365,7 @@ struct HeapTypeGenerator {
   }
 
   HeapTypeKind generateHeapTypeKind() {
-    // Building basic heap types is only allowed in equirecursive mode.
-    // TODO: Relax this.
-    size_t options = 2;
-    if (getTypeSystem() == TypeSystem::Equirecursive) {
-      ++options;
-    }
-    switch (rand.upTo(options)) {
+    switch (rand.upTo(3)) {
       case 0:
         return SignatureKind{};
       case 1:
@@ -441,12 +435,14 @@ struct HeapTypeGenerator {
 
     // Create the heap types.
     for (Index i = 0; i < builder.size(); ++i) {
-      if (!supertypeIndices[i]) {
-        // Create a root type.
-        auto kind = typeKinds[i];
-        if (auto* basic = std::get_if<BasicKind>(&kind)) {
-          builder[i] = *basic;
-        } else if (std::get_if<SignatureKind>(&kind)) {
+      auto kind = typeKinds[i];
+      if (auto* basic = std::get_if<BasicKind>(&kind)) {
+        // The type is already determined.
+        builder[i] = *basic;
+      } else if (!supertypeIndices[i] ||
+                 builder.isBasic(*supertypeIndices[i])) {
+        // No nontrivial supertype, so create a root type.
+        if (std::get_if<SignatureKind>(&kind)) {
           builder[i] = generateSignature();
         } else if (std::get_if<DataKind>(&kind)) {
           if (rand.oneIn(2)) {
@@ -459,12 +455,8 @@ struct HeapTypeGenerator {
         }
       } else {
         // We have a supertype, so create a subtype.
-        Index super = *supertypeIndices[i];
-        HeapType supertype = builder[super];
-        if (builder.isBasic(super)) {
-          auto assignable = generateSubBasic(builder.getBasic(super));
-          std::visit([&](auto&& arg) { builder[i] = arg; }, assignable);
-        } else if (supertype.isSignature()) {
+        HeapType supertype = builder[*supertypeIndices[i]];
+        if (supertype.isSignature()) {
           builder[i] = generateSubSignature(supertype.getSignature());
         } else if (supertype.isStruct()) {
           builder[i] = generateSubStruct(supertype.getStruct());
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 1df14848c..443ab9da2 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -464,16 +464,6 @@ bool isTemp(HeapType type) {
   return !type.isBasic() && getHeapTypeInfo(type)->isTemp;
 }
 
-// Code that traverses the structure of Types often has to be agnostic to the
-// difference between Basic and BasicKind HeapTypes, so uses this helper. On the
-// other hand, canonicalization code often has to differentiate between them so
-// the BasicKind types can be replaced with the corresponding Baic types.
-// BasicKind types should never be visible via the public type API.
-bool isBasicOrBasicKind(HeapType type) {
-  return type.isBasic() ||
-         getHeapTypeInfo(type)->kind == HeapTypeInfo::BasicKind;
-}
-
 // Given a Type that may or may not be backed by the simplest possible
 // representation, return the equivalent type that is definitely backed by the
 // simplest possible representation.
@@ -1811,8 +1801,8 @@ std::ostream& TypePrinter::print(HeapType heapType) {
 #if TRACE_CANONICALIZATION
     os << "[" << ((heapType.getID() >> 4) % 1000) << "]";
     HeapType super;
-    if (heapType.getSuperType(super)) {
-      os << "[super " << ((super.getID() >> 4) % 1000) << "]";
+    if (auto super = heapType.getSuperType()) {
+      os << "[super " << ((super->getID() >> 4) % 1000) << "]";
     }
 #endif
     if (getHeapTypeInfo(heapType)->kind == HeapTypeInfo::BasicKind) {
@@ -2289,7 +2279,6 @@ size_t TypeBuilder::size() { return impl->entries.size(); }
 
 void TypeBuilder::setHeapType(size_t i, HeapType::BasicHeapType basic) {
   assert(i < size() && "Index out of bounds");
-  assert(typeSystem != TypeSystem::Nominal);
   impl->entries[i].set(basic);
 }
 
@@ -2329,9 +2318,6 @@ HeapType TypeBuilder::getTempHeapType(size_t i) {
 }
 
 Type TypeBuilder::getTempTupleType(const Tuple& tuple) {
-  if (typeSystem == TypeSystem::Nominal) {
-    return globalTypeStore.insert(tuple);
-  }
   Type ret = impl->typeStore.insert(tuple);
   if (tuple.types.size() > 1) {
     return markTemp(ret);
@@ -2342,16 +2328,10 @@ Type TypeBuilder::getTempTupleType(const Tuple& tuple) {
 }
 
 Type TypeBuilder::getTempRefType(HeapType type, Nullability nullable) {
-  if (typeSystem == TypeSystem::Nominal) {
-    return globalTypeStore.insert(TypeInfo(type, nullable));
-  }
   return markTemp(impl->typeStore.insert(TypeInfo(type, nullable)));
 }
 
 Type TypeBuilder::getTempRttType(Rtt rtt) {
-  if (typeSystem == TypeSystem::Nominal) {
-    return globalTypeStore.insert(rtt);
-  }
   return markTemp(impl->typeStore.insert(rtt));
 }
 
@@ -2731,7 +2711,7 @@ void ShapeCanonicalizer::initialize(std::vector<HeapType>& roots) {
         TransitionInitializer(Initializer& initializer, size_t parent)
           : initializer(initializer), parent(parent) {}
         void noteChild(HeapType* childType) {
-          if (isBasicOrBasicKind(*childType)) {
+          if (childType->isBasic()) {
             return;
           }
           // Record the transition from parent to child.
@@ -2900,29 +2880,48 @@ void ShapeCanonicalizer::translatePartitionsToTypes() {
 #endif
 }
 
+// Map each temporary Type and HeapType to the locations where they will
+// have to be replaced with canonical Types and HeapTypes.
+struct Locations : TypeGraphWalker<Locations> {
+  std::unordered_map<Type, std::unordered_set<Type*>> types;
+  std::unordered_map<HeapType, std::unordered_set<HeapType*>> heapTypes;
+
+  void preVisitType(Type* type) {
+    if (!type->isBasic()) {
+      types[*type].insert(type);
+    }
+  }
+  void preVisitHeapType(HeapType* ht) {
+    if (!ht->isBasic()) {
+      heapTypes[*ht].insert(ht);
+    }
+  }
+};
+
+void globallyCanonicalizeTypes(Locations& locations) {
+  // Canonicalize non-tuple Types (which never directly refer to other Types)
+  // before tuple Types to avoid canonicalizing a tuple that still contains
+  // non-canonical Types.
+  auto canonicalizeTypes = [&](bool tuples) {
+    for (auto& [original, uses] : locations.types) {
+      if (original.isTuple() == tuples) {
+        Type canonical = globalTypeStore.insert(*getTypeInfo(original));
+        for (Type* use : uses) {
+          *use = canonical;
+        }
+      }
+    }
+  };
+  canonicalizeTypes(false);
+  canonicalizeTypes(true);
+}
+
 // Replaces temporary types and heap types in a type definition graph with their
 // globally canonical versions to prevent temporary types or heap type from
 // leaking into the global stores.
 std::vector<HeapType>
 globallyCanonicalize(std::vector<std::unique_ptr<HeapTypeInfo>>& infos) {
-  // Map each temporary Type and HeapType to the locations where they will
-  // have to be replaced with canonical Types and HeapTypes.
-  struct Locations : TypeGraphWalker<Locations> {
-    std::unordered_map<Type, std::unordered_set<Type*>> types;
-    std::unordered_map<HeapType, std::unordered_set<HeapType*>> heapTypes;
-
-    void preVisitType(Type* type) {
-      if (!type->isBasic()) {
-        types[*type].insert(type);
-      }
-    }
-    void preVisitHeapType(HeapType* ht) {
-      if (!ht->isBasic()) {
-        heapTypes[*ht].insert(ht);
-      }
-    }
-  } locations;
-
+  Locations locations;
   std::vector<HeapType> results;
   results.reserve(infos.size());
   for (auto& info : infos) {
@@ -2954,9 +2953,6 @@ globallyCanonicalize(std::vector<std::unique_ptr<HeapTypeInfo>>& infos) {
   // to become the canonical version. These new canonical HeapTypes still
   // contain references to temporary Types owned by the TypeBuilder, so we must
   // subsequently replace those references with references to canonical Types.
-  // Canonicalize non-tuple Types (which never directly refer to other Types)
-  // before tuple Types to avoid canonicalizing a tuple that still contains
-  // non-canonical Types.
   //
   // Keep a lock on the global HeapType store as long as it can reach temporary
   // types to ensure that no other threads observe the temporary types, for
@@ -2984,20 +2980,7 @@ globallyCanonicalize(std::vector<std::unique_ptr<HeapTypeInfo>>& infos) {
     }
   }
 
-  auto canonicalizeTypes = [&](bool tuples) {
-    for (auto& pair : locations.types) {
-      Type original = pair.first;
-      auto& uses = pair.second;
-      if (original.isTuple() == tuples) {
-        Type canonical = globalTypeStore.insert(*getTypeInfo(original));
-        for (Type* use : uses) {
-          *use = canonical;
-        }
-      }
-    }
-  };
-  canonicalizeTypes(false);
-  canonicalizeTypes(true);
+  globallyCanonicalizeTypes(locations);
 
 #if TRACE_CANONICALIZATION
   std::cerr << "Final Types:\n";
@@ -3010,15 +2993,14 @@ globallyCanonicalize(std::vector<std::unique_ptr<HeapTypeInfo>>& infos) {
   return results;
 }
 
-std::vector<HeapType> buildEquirecursive(TypeBuilder& builder) {
+std::vector<HeapType>
+buildEquirecursive(std::vector<std::unique_ptr<HeapTypeInfo>> infos) {
   std::vector<HeapType> heapTypes;
-  for (auto& entry : builder.impl->entries) {
-    assert(entry.initialized && "Cannot access uninitialized HeapType");
-    entry.info->isFinalized = true;
-    if (!entry.info->isNominal) {
-      entry.info->supertype = nullptr;
+  for (auto& info : infos) {
+    if (!info->isNominal) {
+      info->supertype = nullptr;
     }
-    heapTypes.push_back(entry.get());
+    heapTypes.push_back(asHeapType(info));
   }
 
 #if TIME_CANONICALIZATION
@@ -3128,18 +3110,28 @@ void validateNominalSubTyping(const std::vector<HeapType>& heapTypes) {
   }
 }
 
-std::vector<HeapType> buildNominal(TypeBuilder& builder) {
+std::vector<HeapType>
+buildNominal(std::vector<std::unique_ptr<HeapTypeInfo>> infos) {
 #if TIME_CANONICALIZATION
   auto start = std::chrono::steady_clock::now();
 #endif
 
-  // Just move the HeapTypes to the global store. The Types are already in the
-  // global store.
+  // Move the HeapTypes and the Types they reach to the global stores. First
+  // copy reachable temporary types into the global type store and replace their
+  // uses in the HeapTypeInfos. Then move all the HeapTypeInfos into the global
+  // store. We have to copy types first because correctly hashing the
+  // HeapTypeInfos depends on them reaching only canonical types. That also
+  // means we can't reuse `globallyCanonicalize` here.
+  Locations locations;
+  for (auto& info : infos) {
+    HeapType root = asHeapType(info);
+    locations.walkRoot(&root);
+  }
+  globallyCanonicalizeTypes(locations);
+
   std::vector<HeapType> heapTypes;
-  for (auto& entry : builder.impl->entries) {
-    assert(entry.initialized && "Cannot access uninitialized HeapType");
-    entry.info->isFinalized = true;
-    heapTypes.push_back(globalHeapTypeStore.insert(std::move(entry.info)));
+  for (auto& info : infos) {
+    heapTypes.push_back(globalHeapTypeStore.insert(std::move(info)));
   }
 
 #if TIME_CANONICALIZATION
@@ -3172,16 +3164,95 @@ std::vector<HeapType> buildNominal(TypeBuilder& builder) {
   return heapTypes;
 }
 
+void replaceBasicHeapTypes(std::vector<std::unique_ptr<HeapTypeInfo>>& infos) {
+  // 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
+  // and are never canonical.
+  for (auto& info : infos) {
+    struct BasicTypeReplacer : HeapTypeChildWalker<BasicTypeReplacer> {
+      void noteChild(HeapType* child) {
+        if (child->isBasic()) {
+          // This is already a real basic type. No canonicalization necessary.
+          return;
+        }
+        auto* info = getHeapTypeInfo(*child);
+        if (info->kind == HeapTypeInfo::BasicKind) {
+          *child = info->basic;
+        }
+      }
+    };
+    HeapType type = asHeapType(info);
+    BasicTypeReplacer().walkRoot(&type);
+    if (info->supertype && info->supertype->kind == HeapTypeInfo::BasicKind) {
+      info->supertype = nullptr;
+    }
+  }
+}
+
 } // anonymous namespace
 
 std::vector<HeapType> TypeBuilder::build() {
+  size_t entryCount = impl->entries.size();
+  std::vector<std::optional<HeapType>> basicHeapTypes(entryCount);
+  std::vector<std::unique_ptr<HeapTypeInfo>> toBuild;
+
+  // Mark the entries finalized and "build" basic heap types.
+  for (size_t i = 0; i < entryCount; ++i) {
+    assert(impl->entries[i].initialized &&
+           "Cannot access uninitialized HeapType");
+    auto& info = impl->entries[i].info;
+    info->isFinalized = true;
+    if (info->kind == HeapTypeInfo::BasicKind) {
+      basicHeapTypes[i] = info->basic;
+    } else {
+      toBuild.emplace_back(std::move(info));
+    }
+  }
+
+#if TRACE_CANONICALIZATION
+  auto dumpTypes = [&]() {
+    for (size_t i = 0, j = 0; i < basicHeapTypes.size(); ++i) {
+      if (basicHeapTypes[i]) {
+        std::cerr << i << ": " << *basicHeapTypes[i] << "\n";
+      } else {
+        std::cerr << i << ": " << asHeapType(toBuild[j++]) << "\n";
+      }
+    }
+  };
+  std::cerr << "Before replacing basic heap types:\n";
+  dumpTypes();
+#endif // TRACE_CANONICALIZATION
+
+  // Eagerly replace references to built basic heap types so the more
+  // complicated canonicalization algorithms don't need to consider them.
+  replaceBasicHeapTypes(toBuild);
+
+#if TRACE_CANONICALIZATION
+  std::cerr << "After replacing basic heap types:\n";
+  dumpTypes();
+#endif // TRACE_CANONICALIZATION
+
+  std::vector<HeapType> built;
   switch (typeSystem) {
     case TypeSystem::Equirecursive:
-      return buildEquirecursive(*this);
+      built = buildEquirecursive(std::move(toBuild));
+      break;
     case TypeSystem::Nominal:
-      return buildNominal(*this);
+      built = buildNominal(std::move(toBuild));
+      break;
+  }
+
+  // Combine the basic types with the built types.
+  std::vector<HeapType> result(entryCount);
+  for (size_t i = 0, builtIndex = 0; i < entryCount; ++i) {
+    if (basicHeapTypes[i]) {
+      result[i] = *basicHeapTypes[i];
+    } else {
+      result[i] = built[builtIndex++];
+    }
   }
-  WASM_UNREACHABLE("unexpected type system");
+  return result;
 }
 
 } // namespace wasm