2 files changed, 84 insertions, 14 deletions

diff --git a/src/passes/TypeMerging.cpp b/src/passes/TypeMerging.cpp
index d0877d54f..2ef6f8f54 100644
--- a/src/passes/TypeMerging.cpp
+++ b/src/passes/TypeMerging.cpp
@@ -187,8 +187,20 @@ void TypeMerging::run(Module* module_) {
   // Map each type to its partition in the list.
   std::unordered_map<HeapType, Partitions::iterator> typePartitions;
 
-  // Map the top-level structures of root types to their partitions in the list.
-  std::unordered_map<HeapType, Partitions::iterator, ShapeHash, ShapeEq>
+  // Map optional supertypes and the top-level structures of their refined
+  // children to partitions so that different children that refine the supertype
+  // in the same way can be assigned to the same partition and potentially
+  // merged.
+  // TODO: This is not fully general because it still prevents types from being
+  // merged if they are identical subtypes of two other types that end up being
+  // merged. Fixing this would require 1) merging such input partitions and
+  // re-running DFA minimization until convergence or 2) treating supertypes as
+  // special transitions in the DFA and augmenting Valmari-Lehtinen DFA
+  // minimization so that these special transitions cannot be used to split a
+  // partition if they are self-transitions.
+  std::unordered_map<
+    std::optional<HeapType>,
+    std::unordered_map<HeapType, Partitions::iterator, ShapeHash, ShapeEq>>
     shapePartitions;
 
 #if TYPE_MERGING_DEBUG
@@ -220,9 +232,10 @@ void TypeMerging::run(Module* module_) {
   };
 
   // Similar to the above, but look up or create a partition associated with the
-  // type's top-level shape rather than its identity.
+  // type's supertype and top-level shape rather than its identity.
   auto ensureShapePartition = [&](HeapType type) -> Partitions::iterator {
-    auto [it, inserted] = shapePartitions.insert({type, partitions.end()});
+    auto [it, inserted] =
+      shapePartitions[type.getSuperType()].insert({type, partitions.end()});
     if (inserted) {
       it->second = partitions.insert(partitions.end(), Partition{});
     }
@@ -243,22 +256,16 @@ void TypeMerging::run(Module* module_) {
       ensurePartition(type);
       continue;
     }
+    // If there is no supertype to merge with or if this type refines its
+    // supertype, then we can still potentially merge it with sibling types with
+    // the same structure. Find and add to the partition with other such types.
     auto super = type.getSuperType();
-    // If there is no supertype to merge with, then we can still merge with
-    // other root types with the same structure. Find and add to the partition
-    // with other such types.
-    if (!super) {
+    if (!super || !shapeEq(type, *super)) {
       auto it = ensureShapePartition(type);
       it->push_back(makeDFAState(type));
       typePartitions[type] = it;
       continue;
     }
-    // If this type refines its supertype in some way, then we cannot merge it.
-    // Create a new partition for it.
-    if (!shapeEq(type, *super)) {
-      ensurePartition(type);
-      continue;
-    }
     // The current type and its supertype have the same top-level structure and
     // are not distinguished, so add the current type to its supertype's
     // partition.
diff --git a/test/lit/passes/type-merging.wast b/test/lit/passes/type-merging.wast
index d45bed3c6..ecdc6f961 100644
--- a/test/lit/passes/type-merging.wast
+++ b/test/lit/passes/type-merging.wast
@@ -474,6 +474,69 @@
   )
 )
 
+(module
+  (rec
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $A (struct (field anyref)))
+    (type $A (struct         anyref))
+    ;; CHECK:       (type $B (struct_subtype (field eqref) $A))
+    (type $B (struct_subtype eqref $A))
+    (type $C (struct_subtype eqref $A))
+  )
+
+  ;; CHECK:       (type $none_=>_none (func))
+
+  ;; CHECK:      (func $foo (type $none_=>_none)
+  ;; CHECK-NEXT:  (local $a (ref null $A))
+  ;; CHECK-NEXT:  (local $b (ref null $B))
+  ;; CHECK-NEXT:  (local $c (ref null $B))
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $foo
+    ;; This is the same as above, but now B and C refine A such that they have a
+    ;; different top-level structure. They can still be merged.
+    (local $a (ref null $A))
+    (local $b (ref null $B))
+    (local $c (ref null $C))
+  )
+)
+
+(module
+  (rec
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $A (struct (field anyref)))
+    (type $A (struct         anyref))
+    (type $B (struct_subtype anyref $A))
+    (type $C (struct_subtype anyref $A))
+    ;; CHECK:       (type $E (struct_subtype (field eqref) $A))
+
+    ;; CHECK:       (type $D (struct_subtype (field eqref) $A))
+    (type $D (struct_subtype eqref $B))
+    (type $E (struct_subtype eqref $C))
+  )
+
+  ;; CHECK:       (type $none_=>_none (func))
+
+  ;; CHECK:      (func $foo (type $none_=>_none)
+  ;; CHECK-NEXT:  (local $a (ref null $A))
+  ;; CHECK-NEXT:  (local $b (ref null $A))
+  ;; CHECK-NEXT:  (local $c (ref null $A))
+  ;; CHECK-NEXT:  (local $d (ref null $D))
+  ;; CHECK-NEXT:  (local $e (ref null $E))
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $foo
+    ;; D and E should be mergeable because they have identical shapes and will
+    ;; be siblings after B and C get merged, but we don't support this case yet.
+    ;; TODO: support this.
+    (local $a (ref null $A))
+    (local $b (ref null $B))
+    (local $c (ref null $C))
+    (local $d (ref null $D))
+    (local $e (ref null $E))
+  )
+)
+
 ;; Check that we refinalize properly.
 (module
   ;; CHECK:      (rec