Heap2Local: Properly handle failing array casts (#6772)

Followup to #6727 which added support for failing casts in Struct2Local, but it
turns out that it required Array2Struct changes as well. Specifically, when we
turn an array into a struct then casts can look like they behave differently
(what used to be an array input, becomes a struct), so like with RefTest that we
already handled, check if the cast succeeds in the original form and handle
that.

2 files changed, 164 insertions, 5 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 1e747d6ab..c8d478ad5 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -862,6 +862,11 @@ struct Array2Struct : PostWalker<Array2Struct> {
   // The original type of the allocation, before we turn it into a struct.
   Type originalType;
 
+  // The type of the struct we are changing to (nullable and non-nullable
+  // variations).
+  Type nullStruct;
+  Type nonNullStruct;
+
   Array2Struct(Expression* allocation,
                EscapeAnalyzer& analyzer,
                Function* func,
@@ -928,9 +933,15 @@ struct Array2Struct : PostWalker<Array2Struct> {
     // lowered away to locals anyhow.
     auto nullArray = Type(arrayType, Nullable);
     auto nonNullArray = Type(arrayType, NonNullable);
-    auto nullStruct = Type(structType, Nullable);
-    auto nonNullStruct = Type(structType, NonNullable);
+    nullStruct = Type(structType, Nullable);
+    nonNullStruct = Type(structType, NonNullable);
     for (auto* reached : analyzer.reached) {
+      if (reached->is<RefCast>()) {
+        // Casts must be handled later: We need to see the old type, and to
+        // potentially replace the cast based on that, see below.
+        continue;
+      }
+
       // We must check subtyping here because the allocation may be upcast as it
       // flows around. If we do see such upcasting then we are refining here and
       // must refinalize.
@@ -1032,15 +1043,14 @@ struct Array2Struct : PostWalker<Array2Struct> {
   }
 
   // Some additional operations need special handling
+
   void visitRefTest(RefTest* curr) {
     if (!analyzer.reached.count(curr)) {
       return;
     }
 
     // When we ref.test an array allocation, we cannot simply turn the array
-    // into a struct, as then the test will behave different. (Note that this is
-    // not a problem for ref.*cast*, as the cast simply goes away when the value
-    // flows through, and we verify it will do so in the escape analysis.) To
+    // into a struct, as then the test will behave differently. To properly
     // handle this, check if the test succeeds or not, and write out the outcome
     // here (similar to Struct2Local::visitRefTest). Note that we test on
     // |originalType| here and not |allocation->type|, as the allocation has
@@ -1050,6 +1060,30 @@ struct Array2Struct : PostWalker<Array2Struct> {
                                         builder.makeConst(Literal(result))));
   }
 
+  void visitRefCast(RefCast* curr) {
+    if (!analyzer.reached.count(curr)) {
+      return;
+    }
+
+    // As with RefTest, we need to check if the cast succeeds with the array
+    // type before we turn it into a struct type (as after that change, the
+    // outcome of the cast will look different).
+    if (!Type::isSubType(originalType, curr->type)) {
+      // The cast fails, ensure we trap with an unreachable.
+      replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
+                                          builder.makeUnreachable()));
+    } else {
+      // The cast succeeds. Update the type. (It is ok to use the non-nullable
+      // type here unconditionally, since we know the allocation flows through
+      // here, and anyhow we will be removing the reference during Struct2Local,
+      // later.)
+      curr->type = nonNullStruct;
+    }
+
+    // Regardless of how we altered the type here, refinalize.
+    refinalize = true;
+  }
+
   // Get the value in an expression we know must contain a constant index.
   Index getIndex(Expression* curr) {
     return curr->cast<Const>()->value.getUnsigned();
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 2ce3e12dd..2e38b7e4e 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -4303,3 +4303,128 @@
     )
   )
 )
+
+;; Array casts to structs and arrays.
+(module
+ ;; CHECK:      (type $1 (func))
+
+ ;; CHECK:      (type $0 (func (result (ref struct))))
+ (type $0 (func (result (ref struct))))
+ ;; CHECK:      (type $3 (func (result structref)))
+
+ ;; CHECK:      (type $4 (func (result (ref array))))
+
+ ;; CHECK:      (type $array (array i8))
+ (type $array (array i8))
+
+  ;; CHECK:      (func $array.cast.struct (type $0) (result (ref struct))
+  ;; CHECK-NEXT:  (local $eq (ref eq))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $array.cast.struct (result (ref struct))
+    (local $eq (ref eq))
+    ;; This cast will fail: we cast an array to struct. That we go through
+    ;; (ref eq) in the middle, which seems like it could cast to struct, should
+    ;; not confuse us. And, as the cast fails, the reference does not escape.
+    ;; We can optimize here and will emit an unreachable for the failing cast.
+    (ref.cast (ref struct)
+      (local.tee $eq
+        (array.new_fixed $array 0)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $array.cast.struct.null (type $3) (result structref)
+  ;; CHECK-NEXT:  (local $eq (ref eq))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (unreachable)
+  ;; CHECK-NEXT: )
+  (func $array.cast.struct.null (result (ref null struct))
+    (local $eq (ref eq))
+    ;; As above but the cast is to a nullable type, which changes nothing.
+    (ref.cast (ref null struct)
+      (local.tee $eq
+        (array.new_fixed $array 0)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $array.cast.array (type $4) (result (ref array))
+  ;; CHECK-NEXT:  (local $eq (ref eq))
+  ;; CHECK-NEXT:  (ref.cast (ref array)
+  ;; CHECK-NEXT:   (local.tee $eq
+  ;; CHECK-NEXT:    (array.new_fixed $array 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $array.cast.array (result (ref array))
+    (local $eq (ref eq))
+    ;; Now we cast to array, and the cast succeeds, so we escape, and do
+    ;; nothing to optimize.
+    (ref.cast (ref array)
+      (local.tee $eq
+        (array.new_fixed $array 0)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $array.cast.array.set (type $1)
+  ;; CHECK-NEXT:  (local $eq (ref eq))
+  ;; CHECK-NEXT:  (local $array (ref array))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $array.cast.array.set
+    (local $eq (ref eq))
+    (local $array (ref array))
+    ;; As above, but now we store the result in a local rather than return it
+    ;; out, so it does not escape.
+    (local.set $array
+      (ref.cast (ref array)
+        (local.tee $eq
+          (array.new_fixed $array 0)
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $array.cast.struct.set (type $1)
+  ;; CHECK-NEXT:  (local $eq (ref eq))
+  ;; CHECK-NEXT:  (local $struct (ref struct))
+  ;; CHECK-NEXT:  (local.tee $struct
+  ;; CHECK-NEXT:   (block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (block (result nullref)
+  ;; CHECK-NEXT:      (ref.null none)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (unreachable)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $array.cast.struct.set
+    (local $eq (ref eq))
+    (local $struct (ref struct))
+    ;; As above, but now the cast fails and is stored to a struct local. We do not
+    ;; escape and we emit an unreachable for the cast.
+    (local.set $struct
+      (ref.cast (ref struct)
+        (local.tee $eq
+          (array.new_fixed $array 0)
+        )
+      )
+    )
+  )
+)