[Wasm GC] Allow TypeMerging to merge cast types with TNH (#5425)

If traps are assumed never to happen, then ref.cast and call_indirect
instructions cannot differentiate between types because they always succeed.
Take advantage of this fact by only having these instructions inhibit type
merges when we are not assuming traps never happen.

2 files changed, 132 insertions, 4 deletions

diff --git a/src/passes/TypeMerging.cpp b/src/passes/TypeMerging.cpp
index 1345071d1..aa9fcca67 100644
--- a/src/passes/TypeMerging.cpp
+++ b/src/passes/TypeMerging.cpp
@@ -57,13 +57,26 @@ using ReferredTypes = SmallUnorderedSet<HeapType, 5>;
 struct CastFinder : public PostWalker<CastFinder> {
   ReferredTypes referredTypes;
 
+  // If traps never happen, then ref.cast and call_indirect can never
+  // differentiate between types since they always succeed. Take advantage of
+  // that by not having those instructions inhibit merges in TNH mode.
+  // mode.
+  bool trapsNeverHappen;
+
+  CastFinder(const PassOptions& options)
+    : trapsNeverHappen(options.trapsNeverHappen) {}
+
   template<typename T> void visitCast(T* curr) {
     if (auto type = curr->getCastType(); type != Type::unreachable) {
       referredTypes.insert(type.getHeapType());
     }
   }
 
-  void visitRefCast(RefCast* curr) { visitCast(curr); }
+  void visitRefCast(RefCast* curr) {
+    if (!trapsNeverHappen) {
+      visitCast(curr);
+    }
+  }
 
   void visitRefTest(RefTest* curr) { visitCast(curr); }
 
@@ -74,7 +87,9 @@ struct CastFinder : public PostWalker<CastFinder> {
   }
 
   void visitCallIndirect(CallIndirect* curr) {
-    referredTypes.insert(curr->heapType);
+    if (!trapsNeverHappen) {
+      referredTypes.insert(curr->heapType);
+    }
   }
 };
 
@@ -108,14 +123,14 @@ struct TypeMerging : public Pass {
           return;
         }
 
-        CastFinder finder;
+        CastFinder finder(getPassOptions());
         finder.walk(func->body);
         referredTypes = std::move(finder.referredTypes);
       });
 
     // Also find cast types in the module scope (not possible in the current
     // spec, but do it to be future-proof).
-    CastFinder moduleFinder;
+    CastFinder moduleFinder(getPassOptions());
     moduleFinder.walkModuleCode(module);
 
     // Accumulate all the referredTypes.
diff --git a/test/lit/passes/type-merging-tnh.wast b/test/lit/passes/type-merging-tnh.wast
new file mode 100644
index 000000000..5099760de
--- /dev/null
+++ b/test/lit/passes/type-merging-tnh.wast
@@ -0,0 +1,113 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: foreach %s %t wasm-opt --nominal --closed-world -tnh --type-merging -all -S -o - | filecheck %s
+
+;; ref.cast does not inhibit merging if traps never happen.
+(module
+  ;; CHECK:      (type $A (struct ))
+  (type $A (struct))
+  (type $B (struct_subtype $A))
+
+  ;; CHECK:      (type $ref|$A|_=>_ref|$A| (func (param (ref $A)) (result (ref $A))))
+
+  ;; CHECK:      (func $test (type $ref|$A|_=>_ref|$A|) (param $a (ref $A)) (result (ref $A))
+  ;; CHECK-NEXT:  (ref.cast $A
+  ;; CHECK-NEXT:   (local.get $a)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test (param $a (ref $A)) (result (ref $B))
+    (ref.cast $B
+      (local.get $a)
+    )
+  )
+)
+
+;; Check that a ref.test still inhibits merging with -tnh.
+(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 still inhibits merging with -tnh.
+(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)
+  )
+)
+
+;; call_indirect should not inhibit merging if traps never happen, but we
+;; haven't implemented function type merging yet TODO.
+(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)
+    )
+  )
+)