[Custom Descriptors] Fix TypeMerging topological sort (#7764)

TypeMerging uses a topological sort to process supertypes first in
various places. It also treats descriptor chains as single units
represented by their base described type. Types in WebAssembly cannot
have more than one supertype, but the types in a descriptor chain will
have different supertypes and those supertypes may not necessarily be
all in a single descriptor chain together. As a result, when descriptor
chains are treated as units, it is possible for one chain to have
multiple supertype chains. The topological sort in TypeMerging did not
previously take this into account, and as a result the pass tried to
merge types in the wrong order, leading to assertion failures. Fix the
topological sort to correctly treat descriptor chains as units.

Fixes #7757.

Author: tlively

src/passes/TypeMerging.cpp

@@ -45,8 +45,7 @@
 #include "pass.h"
 #include "support/dfa_minimization.h"
 #include "support/small_set.h"
-#include "wasm-builder.h"
-#include "wasm-type-ordering.h"
+#include "support/topological_sort.h"
 #include "wasm-type.h"
 #include "wasm.h"
 
@@ -169,13 +168,31 @@ struct TypeMerging : public Pass {
 
   std::vector<HeapType>
   mergeableSupertypesFirst(const std::vector<HeapType>& types) {
-    return HeapTypeOrdering::supertypesFirst(
-      types, [&](HeapType type) -> std::optional<HeapType> {
+    // Topological sort so that supertypes come first. Since we treat descriptor
+    // chains as units represented by their base described types, we must handle
+    // the case where one chain has multiple unrelated chains as supertypes.
+    InsertOrderedMap<HeapType, std::vector<HeapType>> subtypes;
+    for (auto type : types) {
+      // Skip descriptor types, since they will be considered as a unit with
+      // their base described types.
+      if (type.getDescribedType()) {
+        continue;
+      }
+      subtypes.insert({type, {}});
+    }
+    // Find the base described type (`superBase`) for each supertype in the
+    // chain starting at `subBase`.
+    for (auto [subBase, _] : subtypes) {
+      for (auto type : subBase.getDescriptorChain()) {
         if (auto super = type.getDeclaredSuperType()) {
-          return getMerged(*super);
+          auto superBase = getMerged(getBaseDescribedType(*super));
+          if (auto it = subtypes.find(superBase); it != subtypes.end()) {
+            it->second.push_back(subBase);
+          }
         }
-        return std::nullopt;
-      });
+      }
+    }
+    return TopologicalSort::sortOf(subtypes.begin(), subtypes.end());
   }
 
   void run(Module* module_) override;
@@ -354,13 +371,7 @@ bool TypeMerging::merge(MergeKind kind) {
   // For each type, either create a new partition or add to its supertype's
   // partition.
   for (auto type : mergeableSupertypesFirst(mergeable)) {
-    // Skip descriptor types. Since types in descriptor chains all have to be
-    // merged into matching descriptor chains together, only the base described
-    // type in each chain is considered, and its DFA state will include the
-    // shape of its entire descriptor chain.
-    if (type.getDescribedType()) {
-      continue;
-    }
+    assert(!type.getDescribedType());
     // We need partitions for any public children of this type since those
     // children will participate in the DFA we're creating. We use the base
     // described type of the child because that's the type that the DFA state

test/lit/passes/type-merging-desc.wast

@@ -306,6 +306,23 @@
   (global $B (ref null $B) (ref.null none))
 )
 
+(module
+  ;; We can merge the $B chain into the $A chain, but we cannot merge in the
+  ;; other direction because that would make $B.desc its own subtype.
+  (rec
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $A (sub (descriptor $A.desc (struct))))
+    (type $A (sub (descriptor $A.desc (struct))))
+    ;; CHECK:       (type $A.desc (sub (describes $A (struct))))
+    (type $A.desc (sub (describes $A (struct))))
+    (type $B (sub (descriptor $B.desc (struct))))
+    (type $B.desc (sub $A.desc (describes $B (struct))))
+  )
+
+  ;; CHECK:      (global $B.desc (ref null $A.desc) (ref.null none))
+  (global $B.desc (ref null $B.desc) (ref.null none))
+)
+
 (module
   ;; CHECK:      (type $public (sub (struct)))
   (type $public (sub (struct)))