[WasmGC] Heap2Local: Optimize RefCast failures (#6727)

Previously we just did not optimize cases where our escape analysis showed an
allocation flowed into a cast that failed. However, after inlining there can be
real-world cases where that happens, even in traps-never-happen mode (if the
cast is behind a conditional branch), so it seems worth optimizing.

2 files changed, 111 insertions, 37 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 720ebba40..1e747d6ab 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -379,11 +379,13 @@ struct EscapeAnalyzer {
       }
 
       void visitRefCast(RefCast* curr) {
-        // As it is our allocation that flows through here, we need to
-        // check that the cast will not trap, so that we can continue
-        // to (hopefully) optimize this allocation.
-        if (Type::isSubType(allocation->type, curr->type)) {
-          escapes = false;
+        // Whether the cast succeeds or fails, it does not escape.
+        escapes = false;
+
+        // If the cast fails then the allocation is fully consumed and does not
+        // flow any further (instead, we trap).
+        if (!Type::isSubType(allocation->type, curr->type)) {
+          fullyConsumes = true;
         }
       }
 
@@ -783,24 +785,22 @@ struct Struct2Local : PostWalker<Struct2Local> {
       return;
     }
 
-    // It is safe to optimize out this RefCast, since we proved it
-    // contains our allocation and we have checked that the type of
-    // the allocation is a subtype of the type of the cast, and so
-    // cannot trap.
-    replaceCurrent(curr->ref);
+    // We know this RefCast receives our allocation, so we can see whether it
+    // succeeds or fails.
+    if (Type::isSubType(allocation->type, curr->type)) {
+      // The cast succeeds, so it is a no-op, and we can skip it, since after we
+      // remove the allocation it will not even be needed for validation.
+      replaceCurrent(curr->ref);
+    } else {
+      // The cast fails, so this must trap.
+      replaceCurrent(builder.makeSequence(builder.makeDrop(curr->ref),
+                                          builder.makeUnreachable()));
+    }
 
-    // We need to refinalize after this, as while we know the cast is not
-    // logically needed - the value flowing through will not be used - we do
-    // need validation to succeed even before other optimizations remove the
-    // code. For example:
-    //
-    //  (block (result $B)
-    //   (ref.cast $B
-    //    (block (result $A)
-    //
-    // Without the cast this does not validate, so we need to refinalize
-    // (which will fix this, as we replace the unused value with a null, so
-    // that type will propagate out).
+    // Either way, we need to refinalize here (we either added an unreachable,
+    // or we replaced a cast with the value being cast, which may have a less-
+    // refined type - it will not be used after we remove the allocation, but we
+    // must still fix that up for validation).
     refinalize = true;
   }
 
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 11f56cae2..f00cb60e2 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -2635,24 +2635,34 @@
   (type $A (sub (struct (field (ref null $A)))))
   ;; CHECK:      (type $1 (func (result anyref)))
 
-  ;; CHECK:      (type $B (sub $A (struct (field (ref $A)))))
-  (type $B (sub $A (struct (field (ref $A)))))
+  ;; CHECK:      (type $B (sub $A (struct (field (ref $A)) (field i32))))
+  (type $B (sub $A (struct (field (ref $A)) (field i32))))
+
+  ;; CHECK:      (type $3 (func (result i32)))
 
   ;; CHECK:      (func $func (type $1) (result anyref)
   ;; CHECK-NEXT:  (local $a (ref $A))
   ;; CHECK-NEXT:  (local $1 (ref $A))
-  ;; CHECK-NEXT:  (local $2 (ref $A))
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 (ref $A))
+  ;; CHECK-NEXT:  (local $4 i32)
   ;; CHECK-NEXT:  (ref.cast (ref $B)
   ;; CHECK-NEXT:   (block (result (ref $A))
   ;; CHECK-NEXT:    (drop
   ;; CHECK-NEXT:     (block (result nullref)
-  ;; CHECK-NEXT:      (local.set $2
+  ;; CHECK-NEXT:      (local.set $3
   ;; CHECK-NEXT:       (struct.new $A
   ;; CHECK-NEXT:        (ref.null none)
   ;; CHECK-NEXT:       )
   ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $4
+  ;; CHECK-NEXT:       (i32.const 1)
+  ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (local.set $1
-  ;; CHECK-NEXT:       (local.get $2)
+  ;; CHECK-NEXT:       (local.get $3)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (local.set $2
+  ;; CHECK-NEXT:       (local.get $4)
   ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
@@ -2675,6 +2685,7 @@
             (struct.new $A
               (ref.null none)
             )
+            (i32.const 1)
           )
         )
       )
@@ -2683,16 +2694,24 @@
 
   ;; CHECK:      (func $cast-success (type $1) (result anyref)
   ;; CHECK-NEXT:  (local $0 (ref $A))
-  ;; CHECK-NEXT:  (local $1 (ref $A))
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 (ref $A))
+  ;; CHECK-NEXT:  (local $3 i32)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result nullref)
-  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:    (local.set $2
   ;; CHECK-NEXT:     (struct.new $A
   ;; CHECK-NEXT:      (ref.null none)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (local.set $0
-  ;; CHECK-NEXT:     (local.get $1)
+  ;; CHECK-NEXT:     (local.get $2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (local.get $3)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (ref.null none)
   ;; CHECK-NEXT:   )
@@ -2706,25 +2725,80 @@
           (struct.new $A
             (ref.null none)
           )
+          (i32.const 1)
         )
       )
     )
   )
   ;; CHECK:      (func $cast-failure (type $1) (result anyref)
-  ;; CHECK-NEXT:  (struct.get $B 0
-  ;; CHECK-NEXT:   (ref.cast (ref $B)
-  ;; CHECK-NEXT:    (struct.new $A
-  ;; CHECK-NEXT:     (struct.new $A
-  ;; CHECK-NEXT:      (ref.null none)
+  ;; CHECK-NEXT:  (local $0 (ref null $A))
+  ;; CHECK-NEXT:  (local $1 (ref null $A))
+  ;; CHECK-NEXT:  (block ;; (replaces unreachable StructGet we can't emit)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (block
+  ;; CHECK-NEXT:     (drop
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $1
+  ;; CHECK-NEXT:        (struct.new $A
+  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $0
+  ;; CHECK-NEXT:        (local.get $1)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
   ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (unreachable)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (unreachable)
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $cast-failure (result anyref)
     (struct.get $B 0
-      ;; The allocated $A arrives here, but the cast will fail, so we do not
-      ;; optimize.
+      ;; The allocated $A arrives here, but the cast will fail. We can remove
+      ;; the allocation and put an unreachable here. (Note that the inner
+      ;; struct.new survives, which would take another cycle to remove.)
+      (ref.cast (ref $B)
+        (struct.new $A
+          (struct.new $A
+            (ref.null none)
+          )
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $cast-failure-nofield (type $3) (result i32)
+  ;; CHECK-NEXT:  (local $0 (ref null $A))
+  ;; CHECK-NEXT:  (local $1 (ref null $A))
+  ;; CHECK-NEXT:  (block ;; (replaces unreachable StructGet we can't emit)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (block
+  ;; CHECK-NEXT:     (drop
+  ;; CHECK-NEXT:      (block (result nullref)
+  ;; CHECK-NEXT:       (local.set $1
+  ;; CHECK-NEXT:        (struct.new $A
+  ;; CHECK-NEXT:         (ref.null none)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (local.set $0
+  ;; CHECK-NEXT:        (local.get $1)
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:       (ref.null none)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (unreachable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (unreachable)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $cast-failure-nofield (result i32)
+    ;; As above, but we read from a field that only exists in $B, despite the
+    ;; allocation that flows here being an $A. We should not error on that.
+    (struct.get $B 1  ;; this changes from 0 to 1
       (ref.cast (ref $B)
         (struct.new $A
           (struct.new $A