[Custom Descriptors] Fix TypeMerging sibling criteria (#7842)

When merging sibling types, we must first partition the types by their
supertypes and shapes. With custom descriptors, we consider full
descriptor chains as single units to be merged. We had updated the shape
partitioning to account for this, but not the supertype partitioning. As
a result, we were incorrectly merging chains with different descriptor
supertypes as long as the base described types had no supertype.

Fix the bug by partitioning not just on the base described type's
supertype, but the sequence of all supertypes in the chain.

Author: tlively

src/passes/TypeMerging.cpp

@@ -248,6 +248,30 @@ struct ShapeHash {
   size_t operator()(const HeapType& type) const { return shapeHash(type); }
 };
 
+// The supertypes of each type in a descriptor chain. Hash and equality-compare
+// them using real identity to identify siblings chains that share the same
+// supertypes.
+using ChainSupers = std::vector<std::optional<HeapType>>;
+
+struct ChainSupersEq {
+  bool operator()(const ChainSupers& a, const ChainSupers& b) const {
+    return a == b;
+  }
+};
+
+struct ChainSupersHash {
+  size_t operator()(const ChainSupers& supers) const {
+    auto digest = wasm::hash(supers.size());
+    for (auto& super : supers) {
+      wasm::rehash(digest, !!super);
+      if (super) {
+        wasm::rehash(digest, *super);
+      }
+    }
+    return digest;
+  }
+};
+
 void TypeMerging::run(Module* module_) {
   module = module_;
 
@@ -339,8 +363,10 @@ bool TypeMerging::merge(MergeKind kind) {
   // that siblings that refine the supertype in the same way can be assigned to
   // the same partition and potentially merged.
   std::unordered_map<
-    std::optional<HeapType>,
-    std::unordered_map<HeapType, Partitions::iterator, ShapeHash, ShapeEq>>
+    std::vector<std::optional<HeapType>>,
+    std::unordered_map<HeapType, Partitions::iterator, ShapeHash, ShapeEq>,
+    ChainSupersHash,
+    ChainSupersEq>
     shapePartitions;
 
   // Ensure the type has a partition and return a reference to it. Since we
@@ -358,12 +384,16 @@ bool TypeMerging::merge(MergeKind kind) {
   // Similar to the above, but look up or create a partition associated with the
   // type's supertype and top-level shape rather than its identity.
   auto ensureShapePartition = [&](HeapType type) -> Partitions::iterator {
-    auto super = type.getDeclaredSuperType();
-    if (super) {
-      super = getMerged(*super);
+    ChainSupers supers;
+    for (auto t : type.getDescriptorChain()) {
+      auto super = t.getDeclaredSuperType();
+      if (super) {
+        super = getMerged(*super);
+      }
+      supers.push_back(super);
     }
     auto [it, inserted] =
-      shapePartitions[super].insert({type, partitions.end()});
+      shapePartitions[supers].insert({type, partitions.end()});
     if (inserted) {
       it->second = partitions.insert(partitions.end(), Partition{});
     }

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

@@ -306,23 +306,6 @@
   (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)))
@@ -368,3 +351,59 @@
   ;; CHECK:      (export "public" (global $public))
   (export "public" (global $public))
 )
+
+(module
+  ;; In principle we could merge the $B chain into the $A chain, because $A and
+  ;; $B are mergeable as siblings and $B.desc is mergeable into its supertype
+  ;; A.desc. However we do not currently do this kind of merge in either phase.
+  (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))))
+    ;; CHECK:       (type $B (sub (descriptor $B.desc (struct))))
+    (type $B (sub (descriptor $B.desc (struct))))
+    ;; CHECK:       (type $B.desc (sub $A.desc (describes $B (struct))))
+    (type $B.desc (sub $A.desc (describes $B (struct))))
+  )
+
+  ;; CHECK:      (global $B.desc (ref null $B.desc) (ref.null none))
+  (global $B.desc (ref null $B.desc) (ref.null none))
+)
+
+(module
+  (rec
+    ;; Like above, but now we have an unrelated C chain. We have these chains:
+    ;;
+    ;; A -> A.desc
+    ;;        ^
+    ;; B -> B.desc
+    ;;
+    ;; C -> C.desc
+    ;;
+    ;; The C chain should not be considered a sibling of the B chain because
+    ;; C.desc is not a subtype of A.desc, even though they would look like
+    ;; siblings if you only considered the base type in the chain.
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $A (descriptor $A.desc (struct)))
+    (type $A (descriptor $A.desc (struct)))
+    ;; CHECK:       (type $A.desc (sub (describes $A (struct))))
+    (type $A.desc (sub (describes $A (struct))))
+    ;; CHECK:       (type $B (descriptor $B.desc (struct)))
+    (type $B (descriptor $B.desc (struct)))
+    ;; CHECK:       (type $B.desc (sub $A.desc (describes $B (struct))))
+    (type $B.desc (sub $A.desc (describes $B (struct))))
+    ;; CHECK:       (type $C (descriptor $C.desc (struct)))
+    (type $C (descriptor $C.desc (struct)))
+    ;; CHECK:       (type $C.desc (describes $C (struct)))
+    (type $C.desc (describes $C (struct)))
+  )
+
+  ;; CHECK:      (global $C (ref null $C) (ref.null none))
+  (global $C (ref null $C) (ref.null none))
+
+  ;; This would become invalid if $B.desc were merged into $C.desc.
+  ;; CHECK:      (global $A.desc (ref null $A.desc) (struct.new_default $B.desc))
+  (global $A.desc (ref null $A.desc) (struct.new_default $B.desc))
+)