[Wasm GC] Fix cast finding in TypeMerging (#5424)

The cast finding code in TypeMerging has been broken since we refactored how
casts work to support nullable casts. The cast finding code only considered
HeapType fields and the return type of `RefCast` to be potential cast targets,
but `RefTest` and `BrOn` instructions now store their cast types as `Type`
fields. Since cast targets are represented more heterogenously now, do not use
the delegations macros to find them.

Add tests showing that each cast instruction independently inhibits merging.

2 files changed, 104 insertions, 27 deletions

diff --git a/src/passes/TypeMerging.cpp b/src/passes/TypeMerging.cpp
index c5028ee20..1345071d1 100644
--- a/src/passes/TypeMerging.cpp
+++ b/src/passes/TypeMerging.cpp
@@ -54,40 +54,27 @@ namespace {
 // Most functions do no casts, or perhaps cast |this| and perhaps a few others.
 using ReferredTypes = SmallUnorderedSet<HeapType, 5>;
 
-struct CastFinder
-  : public PostWalker<CastFinder, UnifiedExpressionVisitor<CastFinder>> {
+struct CastFinder : public PostWalker<CastFinder> {
   ReferredTypes referredTypes;
 
-  void visitExpression(Expression* curr) {
-    // Find all references to a heap type.
-
-#define DELEGATE_ID curr->_id
-
-#define DELEGATE_START(id) [[maybe_unused]] auto* cast = curr->cast<id>();
+  template<typename T> void visitCast(T* curr) {
+    if (auto type = curr->getCastType(); type != Type::unreachable) {
+      referredTypes.insert(type.getHeapType());
+    }
+  }
 
-#define DELEGATE_FIELD_HEAPTYPE(id, field) referredTypes.insert(cast->field);
+  void visitRefCast(RefCast* curr) { visitCast(curr); }
 
-#define DELEGATE_FIELD_CHILD(id, field)
-#define DELEGATE_FIELD_OPTIONAL_CHILD(id, field)
-#define DELEGATE_FIELD_INT(id, field)
-#define DELEGATE_FIELD_INT_ARRAY(id, field)
-#define DELEGATE_FIELD_LITERAL(id, field)
-#define DELEGATE_FIELD_NAME(id, field)
-#define DELEGATE_FIELD_NAME_VECTOR(id, field)
-#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field)
-#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field)
-#define DELEGATE_FIELD_SCOPE_NAME_USE_VECTOR(id, field)
-#define DELEGATE_FIELD_TYPE(id, field)
-#define DELEGATE_FIELD_ADDRESS(id, field)
-#define DELEGATE_FIELD_CHILD_VECTOR(id, field)
+  void visitRefTest(RefTest* curr) { visitCast(curr); }
 
-#include "wasm-delegations-fields.def"
+  void visitBrOn(BrOn* curr) {
+    if (curr->op == BrOnCast || curr->op == BrOnCastFail) {
+      visitCast(curr);
+    }
   }
 
-  void visitRefCast(Expression* curr) {
-    if (curr->type != Type::unreachable) {
-      referredTypes.insert(curr->type.getHeapType());
-    }
+  void visitCallIndirect(CallIndirect* curr) {
+    referredTypes.insert(curr->heapType);
   }
 };
 
diff --git a/test/lit/passes/type-merging.wast b/test/lit/passes/type-merging.wast
index 85565551d..b4b84f94a 100644
--- a/test/lit/passes/type-merging.wast
+++ b/test/lit/passes/type-merging.wast
@@ -276,3 +276,93 @@
     (local $c (ref null $sub-refarray-nn))
   )
 )
+
+;; Check that a ref.test inhibits merging (ref.cast is already checked above).
+(module
+  ;; CHECK:      (type $ref|$A|_=>_i32 (func (param (ref $A)) (result i32)))
+
+  ;; CHECK:      (type $A (struct ))
+  (type $A (struct))
+  ;; CHECK:      (type $B (struct_subtype  $A))
+  (type $B (struct_subtype $A))
+
+  ;; CHECK:      (func $test (type $ref|$A|_=>_i32) (param $a (ref $A)) (result i32)
+  ;; CHECK-NEXT:  (ref.test $B
+  ;; CHECK-NEXT:   (local.get $a)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test (param $a (ref $A)) (result i32)
+    (ref.test $B
+      (local.get $a)
+    )
+  )
+)
+
+;; Check that a br_on_cast inhibits merging.
+(module
+  ;; CHECK:      (type $A (struct ))
+  (type $A (struct))
+  ;; CHECK:      (type $B (struct_subtype  $A))
+  (type $B (struct_subtype $A))
+
+  ;; CHECK:      (type $ref|$A|_=>_ref|$B| (func (param (ref $A)) (result (ref $B))))
+
+  ;; CHECK:      (func $test (type $ref|$A|_=>_ref|$B|) (param $a (ref $A)) (result (ref $B))
+  ;; CHECK-NEXT:  (block $__binaryen_fake_return (result (ref $B))
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (br_on_cast $__binaryen_fake_return $B
+  ;; CHECK-NEXT:     (local.get $a)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (block $l (result (ref $A))
+  ;; CHECK-NEXT:     (br_on_non_null $l
+  ;; CHECK-NEXT:      (local.get $a)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (unreachable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (unreachable)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test (param $a (ref $A)) (result (ref $B))
+    (drop
+      (br_on_cast 0 $B
+        (local.get $a)
+      )
+    )
+    ;; Also check that a different br_on* doesn't confuse us.
+    (drop
+      (block $l (result (ref $A))
+        (br_on_non_null $l
+          (local.get $a)
+        )
+        (unreachable)
+      )
+    )
+    (unreachable)
+  )
+)
+
+;; Check that a call_indirect inhibits merging.
+(module
+  ;; CHECK:      (type $A (func))
+  (type $A (func))
+  ;; CHECK:      (type $B (func_subtype $A))
+  (type $B (func_subtype $A))
+
+  (table 1 1 (ref null $A))
+
+  ;; CHECK:      (table $0 1 1 (ref null $A))
+
+  ;; CHECK:      (func $test (type $A)
+  ;; CHECK-NEXT:  (call_indirect $0 (type $B)
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test (type $A)
+    (call_indirect (type $B)
+      (i32.const 0)
+    )
+  )
+)