GUFA: Refine casts (#5805)

If we see (ref.cast $A) but we have inferred that a more refined type will
be present there at runtime $B then we can refine the cast to (ref.cast $B).

We could do the same even when a cast is not present, but that would
increase code size. This optimization keeps code size constant.

4 files changed, 70 insertions, 8 deletions

diff --git a/src/passes/GUFA.cpp b/src/passes/GUFA.cpp
index 137ffda17..8f58eaf8a 100644
--- a/src/passes/GUFA.cpp
+++ b/src/passes/GUFA.cpp
@@ -255,6 +255,25 @@ struct GUFAOptimizer
     }
   }
 
+  void visitRefCast(RefCast* curr) {
+    auto currType = curr->type;
+    auto inferredType = getContents(curr).getType();
+    if (inferredType.isRef() && inferredType != currType &&
+        Type::isSubType(inferredType, currType)) {
+      // We have inferred that this will only contain something of a more
+      // refined type, so we might as well cast to that more refined type.
+      //
+      // Note that we could in principle apply this in all expressions by adding
+      // a cast. However, to be careful with code size, we only refine existing
+      // casts for now.
+      curr->type = inferredType;
+    }
+
+    // Apply the usual optimizations as well, such as potentially replacing this
+    // with a constant.
+    visitExpression(curr);
+  }
+
   // TODO: If an instruction would trap on null, like struct.get, we could
   //       remove it here if it has no possible contents and if we are in
   //       traps-never-happen mode (that is, we'd have proven it can only trap,
diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 91298106d..fccbc3321 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -191,7 +191,9 @@ struct Heap2LocalOptimizer {
     localGraph.computeSetInfluences();
 
     // All the allocations in the function.
-    // TODO: Arrays (of constant size) as well.
+    // TODO: Arrays (of constant size) as well, if all element accesses use
+    //       constant indexes. One option might be to first convert such
+    //       nonescaping arrays into structs.
     FindAll<StructNew> allocations(func->body);
 
     for (auto* allocation : allocations.list) {
diff --git a/test/lit/passes/gufa-refs.wast b/test/lit/passes/gufa-refs.wast
index 928796823..50658ef34 100644
--- a/test/lit/passes/gufa-refs.wast
+++ b/test/lit/passes/gufa-refs.wast
@@ -2554,7 +2554,7 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result nullref)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (ref.cast null $struct
+  ;; CHECK-NEXT:     (ref.cast null none
   ;; CHECK-NEXT:      (select (result i31ref)
   ;; CHECK-NEXT:       (ref.null none)
   ;; CHECK-NEXT:       (i31.new
@@ -2593,7 +2593,9 @@
     )
     ;; A null or an i31 will reach the cast; only the null can actually pass
     ;; through (an i31 would fail the cast). Given that, we can infer a null for
-    ;; the value of the cast.
+    ;; the value of the cast. (The cast itself will also be turned into a cast
+    ;; to null, but it is dropped right before we return a null, so that has no
+    ;; benefit in this case.)
     (drop
       (ref.cast null $struct
         (select
@@ -5642,3 +5644,38 @@
     )
   )
 )
+
+(module
+  ;; CHECK:      (type $A (struct ))
+  (type $A (struct))
+
+  ;; CHECK:      (type $B (sub $A (struct )))
+  (type $B (sub $A (struct)))
+
+  ;; CHECK:      (type $none_=>_ref|$A| (func (result (ref $A))))
+
+  ;; CHECK:      (type $none_=>_anyref (func (result anyref)))
+
+  ;; CHECK:      (export "func" (func $func))
+
+  ;; CHECK:      (func $func (type $none_=>_ref|$A|) (result (ref $A))
+  ;; CHECK-NEXT:  (ref.cast $B
+  ;; CHECK-NEXT:   (call $get-B-def-any)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $func (export "func") (result (ref $A))
+    ;; Call a function that actually returns a B, though it is defined as
+    ;; returning an anyref. Then cast it to A. We can infer that it will be a B,
+    ;; so we can cast to B here instead.
+    (ref.cast $A
+      (call $get-B-def-any)
+    )
+  )
+
+  ;; CHECK:      (func $get-B-def-any (type $none_=>_anyref) (result anyref)
+  ;; CHECK-NEXT:  (struct.new_default $B)
+  ;; CHECK-NEXT: )
+  (func $get-B-def-any (result anyref)
+    (struct.new $B)
+  )
+)
diff --git a/test/lit/passes/gufa-vs-cfp.wast b/test/lit/passes/gufa-vs-cfp.wast
index 2171a9fc1..0db24bbad 100644
--- a/test/lit/passes/gufa-vs-cfp.wast
+++ b/test/lit/passes/gufa-vs-cfp.wast
@@ -876,7 +876,7 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (struct.get $substruct 0
-  ;; CHECK-NEXT:    (ref.cast null $substruct
+  ;; CHECK-NEXT:    (ref.cast $substruct
   ;; CHECK-NEXT:     (local.get $ref)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
@@ -904,6 +904,8 @@
     )
     (drop
       (struct.get $substruct 0
+        ;; This cast will be refined to be non-nullable, as the LocalGraph
+        ;; analysis will show that it must be so.
         (ref.cast null $substruct
           (local.get $ref)
         )
@@ -949,10 +951,8 @@
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (block (result i32)
   ;; CHECK-NEXT:    (drop
-  ;; CHECK-NEXT:     (struct.get $substruct 0
-  ;; CHECK-NEXT:      (ref.cast null $substruct
-  ;; CHECK-NEXT:       (local.get $ref)
-  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     (ref.cast $substruct
+  ;; CHECK-NEXT:      (local.get $ref)
   ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (i32.const 20)
@@ -981,6 +981,10 @@
     )
     (drop
       (struct.get $substruct 0
+        ;; This cast will be refined to be non-nullable, as the LocalGraph
+        ;; analysis will show that it must be so. After that, the dropped
+        ;; struct.get can be removed as it has no side effects (the only
+        ;; possible effect was a trap on null).
         (ref.cast null $substruct
           (local.get $ref)
         )