5 files changed, 77 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 {