4 files changed, 72 insertions, 5 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index 67ec4a5c0..fa2eec02b 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -36,7 +36,22 @@ void GlobalTypeRewriter::update() {
   // come before their subtypes.
   Index i = 0;
   auto privateTypes = ModuleUtils::getPrivateHeapTypes(wasm);
-  for (auto type : HeapTypeOrdering::SupertypesFirst(privateTypes)) {
+
+  // Topological sort to have supertypes first, but we have to account for the
+  // fact that we may be replacing the supertypes to get the order correct.
+  struct SupertypesFirst
+    : HeapTypeOrdering::SupertypesFirstBase<SupertypesFirst> {
+    GlobalTypeRewriter& parent;
+
+    SupertypesFirst(GlobalTypeRewriter& parent,
+                    const std::vector<HeapType>& types)
+      : SupertypesFirstBase(types), parent(parent) {}
+    std::optional<HeapType> getSuperType(HeapType type) {
+      return parent.getSuperType(type);
+    }
+  };
+
+  for (auto type : SupertypesFirst(*this, privateTypes)) {
     typeIndices[type] = i++;
   }
 
diff --git a/src/passes/TypeMerging.cpp b/src/passes/TypeMerging.cpp
index 5623ecdb9..d0877d54f 100644
--- a/src/passes/TypeMerging.cpp
+++ b/src/passes/TypeMerging.cpp
@@ -292,6 +292,21 @@ void TypeMerging::run(Module* module_) {
     }
   }
 
+#if TYPE_MERGING_DEBUG
+  std::cerr << "Merges):\n";
+  std::unordered_map<HeapType, std::vector<HeapType>> mergees;
+  for (auto& [mergee, target] : merges) {
+    mergees[target].push_back(mergee);
+  }
+  for (auto& [target, types] : mergees) {
+    std::cerr << "target: " << print(target) << "\n";
+    for (auto type : types) {
+      std::cerr << "  " << print(type) << "\n";
+    }
+    std::cerr << "\n";
+  }
+#endif // TYPE_MERGING_DEBUG
+
   applyMerges(merges);
 }
 
diff --git a/src/wasm-type-ordering.h b/src/wasm-type-ordering.h
index 0b0042ca3..af1995de4 100644
--- a/src/wasm-type-ordering.h
+++ b/src/wasm-type-ordering.h
@@ -28,13 +28,14 @@ namespace wasm::HeapTypeOrdering {
 // Given a collection of types, iterate through it such that each type in the
 // collection is visited only after its immediate supertype in the collection is
 // visited.
-template<typename T>
-struct SupertypesFirst : TopologicalSort<HeapType, SupertypesFirst<T>> {
+template<typename SupertypeProvider>
+struct SupertypesFirstBase
+  : TopologicalSort<HeapType, SupertypesFirstBase<SupertypeProvider>> {
   // For each type in the input collection, whether it is a supertype. Used to
   // track membership in the input collection.
   InsertOrderedMap<HeapType, bool> typeSet;
 
-  SupertypesFirst(const T& types) {
+  template<typename T> SupertypesFirstBase(const T& types) {
     for (auto type : types) {
       typeSet[type] = false;
     }
@@ -56,12 +57,22 @@ struct SupertypesFirst : TopologicalSort<HeapType, SupertypesFirst<T>> {
 
   void pushPredecessors(HeapType type) {
     // Do not visit types that weren't in the input collection.
-    if (auto super = type.getSuperType(); super && typeSet.count(*super)) {
+    if (auto super = static_cast<SupertypeProvider*>(this)->getSuperType(type);
+        super && typeSet.count(*super)) {
       this->push(*super);
     }
   }
 };
 
+struct SupertypesFirst : SupertypesFirstBase<SupertypesFirst> {
+  template<typename T>
+  SupertypesFirst(const T& types) : SupertypesFirstBase(types) {}
+
+  std::optional<HeapType> getSuperType(HeapType type) {
+    return type.getSuperType();
+  }
+};
+
 } // namespace wasm::HeapTypeOrdering
 
 #endif // wasm_wasm_type_ordering_h
diff --git a/test/lit/passes/type-merging.wast b/test/lit/passes/type-merging.wast
index 4a58ca2d0..d45bed3c6 100644
--- a/test/lit/passes/type-merging.wast
+++ b/test/lit/passes/type-merging.wast
@@ -690,6 +690,32 @@
   )
 )
 
+;; Regresssion test for a bug in which we merged A into A', but
+;; type-updating.cpp ordered B before A', so the supertype ordering was
+;; incorrect.
+(module
+ (rec
+  (type $A (struct))
+  (type $B (struct_subtype $A))
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $X (struct (field (ref $A'))))
+  (type $X (struct (ref $B)))
+  ;; CHECK:       (type $A' (struct ))
+  (type $A' (struct))
+ )
+ ;; CHECK:       (type $none_=>_none (func))
+
+ ;; CHECK:      (func $foo (type $none_=>_none)
+ ;; CHECK-NEXT:  (local $b (ref null $A'))
+ ;; CHECK-NEXT:  (local $x (ref null $X))
+ ;; CHECK-NEXT:  (nop)
+ ;; CHECK-NEXT: )
+ (func $foo
+   (local $b (ref null $A'))
+   (local $x (ref null $X))
+ )
+)
+
 ;; Check that a ref.test inhibits merging (ref.cast is already checked above).
 (module
   ;; CHECK:      (rec