Fix a fuzz bug in TypeMapper (#5851)

TypeMapper is a utility used to globally rewrite types, mapping some eliminated
source types into destination types they should be replaced with. This was
previously done by first rewriting all the types in the IR according to the
given mapping, then rewriting the type definitions and updating all the types in
the IR again. Not only was doing the rewriting twice inefficient, it also
introduced a subtle bug where the set of private types eligible to be rewritten
could be inconsistent because updating types in the IR could change the types of
control flow structures. The fuzzer found a case where this inconsistency caused
the type rebuilding to fail.

Fix the bug by first building the new types with the mapping applied and only
then rewriting the IR a single time.

Also add a `TypeBuilder::dump` utility for use in debugging.

Fixes #5845.

6 files changed, 112 insertions, 18 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index e11705812..85ead1e9e 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -28,7 +28,9 @@ namespace wasm {
 
 GlobalTypeRewriter::GlobalTypeRewriter(Module& wasm) : wasm(wasm) {}
 
-void GlobalTypeRewriter::update() {
+void GlobalTypeRewriter::update() { mapTypes(rebuildTypes()); }
+
+GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes() {
   // Find the heap types that are not publicly observable. Even in a closed
   // world scenario, don't modify public types because we assume that they may
   // be reflected on or used for linking. Figure out where each private type
@@ -56,7 +58,7 @@ void GlobalTypeRewriter::update() {
   }
 
   if (typeIndices.size() == 0) {
-    return;
+    return {};
   }
   typeBuilder.grow(typeIndices.size());
 
@@ -139,7 +141,7 @@ void GlobalTypeRewriter::update() {
     }
   }
 
-  mapTypes(oldToNewTypes);
+  return oldToNewTypes;
 }
 
 void GlobalTypeRewriter::mapTypes(const TypeMap& oldToNewTypes) {
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index 160934da1..31fd77c69 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -414,6 +414,12 @@ public:
     } rewriter(wasm, updates);
   }
 
+protected:
+  // Builds new types after updating their contents using the hooks below and
+  // returns a map from the old types to the modified types. Used internally in
+  // update().
+  TypeMap rebuildTypes();
+
 private:
   TypeBuilder typeBuilder;
 
@@ -434,13 +440,26 @@ public:
     : GlobalTypeRewriter(wasm), mapping(mapping) {}
 
   void map() {
-    // Map the types of expressions (curr->type, etc.) to their merged
-    // types.
-    mapTypes(mapping);
-
     // Update the internals of types (struct fields, signatures, etc.) to
     // refer to the merged types.
-    update();
+    auto newMapping = rebuildTypes();
+
+    // Compose the user-provided mapping from old types to other old types with
+    // the new mapping from old types to new types. `newMapping` will become
+    // a copy of `mapping` except that the destination types will be the newly
+    // built types.
+    for (auto& [src, dest] : mapping) {
+      if (auto it = newMapping.find(dest); it != newMapping.end()) {
+        newMapping[src] = it->second;
+      } else {
+        // This mapping was to a type that was not rebuilt, perhaps because it
+        // is a basic type. Just use this mapping unmodified.
+        newMapping[src] = dest;
+      }
+    }
+
+    // Map the types of expressions (curr->type, etc.) to the correct new types.
+    mapTypes(newMapping);
   }
 
   Type getNewType(Type type) {
diff --git a/src/passes/TypeMerging.cpp b/src/passes/TypeMerging.cpp
index fa36eb767..808076f85 100644
--- a/src/passes/TypeMerging.cpp
+++ b/src/passes/TypeMerging.cpp
@@ -41,7 +41,6 @@
 #include "pass.h"
 #include "support/dfa_minimization.h"
 #include "support/small_set.h"
-#include "support/topological_sort.h"
 #include "wasm-builder.h"
 #include "wasm-type-ordering.h"
 #include "wasm.h"
@@ -246,8 +245,9 @@ bool TypeMerging::merge(MergeKind kind) {
   Partitions partitions;
 
 #if TYPE_MERGING_DEBUG
+  auto printedPrivateTypes = ModuleUtils::getPrivateHeapTypes(*module);
   using Fallback = IndexedTypeNameGenerator<DefaultTypeNameGenerator>;
-  Fallback printPrivate(ModuleUtils::getPrivateHeapTypes(*module), "private.");
+  Fallback printPrivate(printedPrivateTypes, "private.");
   ModuleTypeNameGenerator<Fallback> print(*module, printPrivate);
   auto dumpPartitions = [&]() {
     size_t i = 0;
@@ -342,7 +342,16 @@ bool TypeMerging::merge(MergeKind kind) {
   }
 
 #if TYPE_MERGING_DEBUG
-  std::cerr << "Initial partitions:\n";
+  std::cerr << "Initial partitions (";
+  switch (kind) {
+    case Supertypes:
+      std::cerr << "supertypes";
+      break;
+    case Siblings:
+      std::cerr << "siblings";
+      break;
+  }
+  std::cerr << "):\n";
   dumpPartitions();
 #endif
 
diff --git a/src/wasm-type.h b/src/wasm-type.h
index 323ef0207..580c198e3 100644
--- a/src/wasm-type.h
+++ b/src/wasm-type.h
@@ -655,6 +655,8 @@ struct TypeBuilder {
   };
 
   Entry operator[](size_t i) { return Entry{*this, i}; }
+
+  void dump();
 };
 
 std::ostream& operator<<(std::ostream&, Type);
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 6baf0feb0..114a83bd6 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -1621,6 +1621,9 @@ void TypePrinter::printHeapTypeName(HeapType type) {
     return;
   }
   os << '$' << generator(type).name;
+#if TRACE_CANONICALIZATION
+  os << "(;" << ((type.getID() >> 4) % 1000) << ";) ";
+#endif
 }
 
 std::ostream& TypePrinter::print(Type type) {
@@ -1643,12 +1646,12 @@ std::ostream& TypePrinter::print(Type type) {
     }
   }
 
-  if (isTemp(type)) {
-    os << "(; temp ;) ";
-  }
 #if TRACE_CANONICALIZATION
   os << "(;" << ((type.getID() >> 4) % 1000) << ";) ";
 #endif
+  if (isTemp(type)) {
+    os << "(; temp ;) ";
+  }
   if (type.isTuple()) {
     print(type.getTuple());
   } else if (type.isRef()) {
@@ -1742,10 +1745,6 @@ std::ostream& TypePrinter::print(HeapType type) {
     os << "(; temp ;) ";
   }
 
-#if TRACE_CANONICALIZATION
-  os << "(;" << ((type.getID() >> 4) % 1000) << ";)";
-#endif
-
   // TODO: Use shorthand for final types once we parse MVP signatures as final.
   bool useSub = false;
   auto super = type.getSuperType();
@@ -2523,6 +2522,34 @@ TypeBuilder::BuildResult TypeBuilder::build() {
   return {results};
 }
 
+void TypeBuilder::dump() {
+  std::vector<HeapType> types;
+  for (size_t i = 0; i < size(); ++i) {
+    types.push_back((*this)[i]);
+  }
+  IndexedTypeNameGenerator<DefaultTypeNameGenerator> print(types);
+
+  std::optional<RecGroup> currGroup;
+  for (auto type : types) {
+    if (auto newGroup = type.getRecGroup(); newGroup != currGroup) {
+      if (currGroup && currGroup->size() > 1) {
+        std::cerr << ")\n";
+      }
+      if (newGroup.size() > 1) {
+        std::cerr << "(rec\n";
+      }
+      currGroup = newGroup;
+    }
+    if (currGroup->size() > 1) {
+      std::cerr << "  ";
+    }
+    std::cerr << print(type) << "\n";
+  }
+  if (currGroup && currGroup->size() > 1) {
+    std::cerr << ")\n";
+  }
+}
+
 } // namespace wasm
 
 namespace std {
diff --git a/test/lit/passes/type-merging.wast b/test/lit/passes/type-merging.wast
index 37d5337ac..66d1e1222 100644
--- a/test/lit/passes/type-merging.wast
+++ b/test/lit/passes/type-merging.wast
@@ -918,6 +918,41 @@
   )
 )
 
+;; Regression test for a bug where we updated module types before building the
+;; new types, causing the set of private types to change unexpectedly and
+;; leading to a failure to build new types.
+;; TODO: Store a heap type on control flow structures to avoid creating
+;; standalone function types for them.
+;; TODO: Investigate why the rec group contains two of the same type below.
+(module
+ (rec
+  (type $A (func (result (ref any) (ref $C))))
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $B (func))
+  (type $B (func))
+  (type $C (sub $B (func)))
+  ;; CHECK:       (type $none_=>_ref|any|_ref|$B| (func (result (ref any) (ref $B))))
+
+  ;; CHECK:       (type $none_=>_ref|any|_ref|$B| (func (result (ref any) (ref $B))))
+
+  ;; CHECK:       (type $D (sub final $none_=>_ref|any|_ref|$B| (func (result (ref any) (ref $B)))))
+  (type $D (sub final $A (func (result (ref any) (ref $C)))))
+ )
+
+ ;; CHECK:      (type $none_=>_ref|any|_ref|$B| (func (result (ref any) (ref $B))))
+
+ ;; CHECK:      (func $test (type $D) (result (ref any) (ref $B))
+ ;; CHECK-NEXT:  (block $l (result (ref any) (ref $B))
+ ;; CHECK-NEXT:   (unreachable)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $test (type $D) (result (ref any) (ref $C))
+  (block $l (result (ref any) (ref $C))
+   (unreachable)
+  )
+ )
+)
+
 ;; Check that a ref.test inhibits merging (ref.cast is already checked above).
 (module
   ;; CHECK:      (rec