[Wasm GC] Fix supertype ordering bug in GlobalTypeRewriter (#5442)

GlobalTypeRewriter made sure to order supertypes before their subtypes so that
type building would succeed, but this ordering previously did not account for
the fact that supertypes might be replaced by the type updating. When a
supertype is replaced, the replacement might not previously have had a
supertype/subtype relationship with the subtype, so it might have happened to be
ordered after the subtype.

Fix the problem by taking supertype replacements into account when ordering
types in the GlobalTypeRewriter. This requires generalizing the
`SupertypesFirst` utility in wasm-type-ordering.h to allow users to modify how
supertypes are found.

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