Make heap2local work through casts (#5952)

E.g.

(local $x (ref eq)
...
(local.set $x
  (struct.new $float
    ...
  )
)
(struct.get $float 0
  (ref.cast (ref $float)
    (local.get $x)
  )
)

This PR allows us to use heap2local, ignoring the passing cast.

This is similar to existing handling of ref.as_non_null.

2 files changed, 122 insertions, 9 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 84037e0dd..e027f5a46 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -430,6 +430,18 @@ struct Heap2LocalOptimizer {
       replaceCurrent(curr->value);
     }
 
+    void visitRefCast(RefCast* curr) {
+      if (!reached.count(curr)) {
+        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);
+    }
+
     void visitStructSet(StructSet* curr) {
       if (!reached.count(curr)) {
         return;
@@ -526,7 +538,7 @@ struct Heap2LocalOptimizer {
         return;
       }
 
-      switch (getParentChildInteraction(parent, child)) {
+      switch (getParentChildInteraction(allocation, parent, child)) {
         case ParentChildInteraction::Escapes: {
           // If the parent may let us escape then we are done.
           return;
@@ -586,7 +598,8 @@ struct Heap2LocalOptimizer {
     rewriter.applyOptimization();
   }
 
-  ParentChildInteraction getParentChildInteraction(Expression* parent,
+  ParentChildInteraction getParentChildInteraction(StructNew* allocation,
+                                                   Expression* parent,
                                                    Expression* child) {
     // If there is no parent then we are the body of the function, and that
     // means we escape by flowing to the caller.
@@ -595,6 +608,7 @@ struct Heap2LocalOptimizer {
     }
 
     struct Checker : public Visitor<Checker> {
+      StructNew* allocation;
       Expression* child;
 
       // Assume escaping (or some other problem we cannot analyze) unless we are
@@ -647,6 +661,15 @@ struct Heap2LocalOptimizer {
         }
       }
 
+      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;
+        }
+      }
+
       // GC operations.
       void visitStructSet(StructSet* curr) {
         // The reference does not escape (but the value is stored to memory and
@@ -664,6 +687,7 @@ struct Heap2LocalOptimizer {
       // TODO Array and I31 operations
     } checker;
 
+    checker.allocation = allocation;
     checker.child = child;
     checker.visit(parent);
 
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 6fa0e2a1f..b814c4987 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -629,7 +629,7 @@
     )
   )
 
-  ;; CHECK:      (func $local-copies-conditional (type $7) (param $x i32) (result f64)
+  ;; CHECK:      (func $local-copies-conditional (type $8) (param $x i32) (result f64)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
@@ -719,7 +719,7 @@
     )
   )
 
-  ;; CHECK:      (func $non-exclusive-get (type $7) (param $x i32) (result f64)
+  ;; CHECK:      (func $non-exclusive-get (type $8) (param $x i32) (result f64)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local.set $ref
   ;; CHECK-NEXT:   (struct.new_default $struct.A)
@@ -751,7 +751,7 @@
     )
   )
 
-  ;; CHECK:      (func $tee (type $8) (result i32)
+  ;; CHECK:      (func $tee (type $5) (result i32)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
@@ -871,7 +871,7 @@
     )
   )
 
-  ;; CHECK:      (func $escape-flow-out (type $5) (result anyref)
+  ;; CHECK:      (func $escape-flow-out (type $6) (result anyref)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (struct.set $struct.A 0
   ;; CHECK-NEXT:   (local.tee $ref
@@ -893,7 +893,7 @@
     (local.get $ref)
   )
 
-  ;; CHECK:      (func $escape-return (type $5) (result anyref)
+  ;; CHECK:      (func $escape-return (type $6) (result anyref)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (struct.set $struct.A 0
   ;; CHECK-NEXT:   (local.tee $ref
@@ -1636,7 +1636,7 @@
     )
   )
 
-  ;; CHECK:      (func $ref-as-non-null-through-local (type $8) (result i32)
+  ;; CHECK:      (func $ref-as-non-null-through-local (type $5) (result i32)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
@@ -1774,7 +1774,7 @@
     )
   )
 
-  ;; CHECK:      (func $non-nullable-local (type $5) (result anyref)
+  ;; CHECK:      (func $non-nullable-local (type $6) (result anyref)
   ;; CHECK-NEXT:  (local $0 (ref null $struct.A))
   ;; CHECK-NEXT:  (local $1 i32)
   ;; CHECK-NEXT:  (local $2 f64)
@@ -1895,6 +1895,41 @@
       (br $loop)
     )
   )
+
+  ;; CHECK:      (func $ref-cast (type $5) (result i32)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local $1 f64)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 f64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (f64.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $0
+  ;; 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:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (local.get $0)
+  ;; CHECK-NEXT: )
+  (func $ref-cast (result i32)
+    (struct.get $struct.A 0
+      (ref.cast (ref $struct.A)
+        (struct.new $struct.A
+          (i32.const 0)
+          (f64.const 0)
+        )
+      )
+    )
+  )
 )
 
 (module
@@ -1945,4 +1980,58 @@
       )
     )
   )
+
+  ;; CHECK:      (func $cast-success (type $1) (result anyref)
+  ;; CHECK-NEXT:  (local $0 (ref $A))
+  ;; CHECK-NEXT:  (local $1 (ref $A))
+  ;; 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:  (local.get $0)
+  ;; CHECK-NEXT: )
+  (func $cast-success (result anyref)
+    (struct.get $A 0
+      (ref.cast (ref $A)
+        (struct.new $B
+          (struct.new $A
+            (ref.null none)
+          )
+        )
+      )
+    )
+  )
+  ;; 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:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; 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.
+      (ref.cast (ref $B)
+        (struct.new $A
+          (struct.new $A
+            (ref.null none)
+          )
+        )
+      )
+    )
+  )
 )