[Wasm GC] Fix optimizations on ref.cast of null, and optimize to ref.as_non_null (#5398)

We were checking the heap type, but now casts care about the nullability as
well.

If the nullability is the only problem, that is, the heap type will be fine but we
might have a null, we can at least switch a ref.cast (non-null) to a
ref.as_non_null.

2 files changed, 82 insertions, 1 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 916991ecb..ed99705ae 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -1963,7 +1963,7 @@ struct OptimizeInstructions
     }
 
     // Check whether the cast will definitely succeed.
-    if (HeapType::isSubType(curr->ref->type.getHeapType(), intendedType)) {
+    if (Type::isSubType(curr->ref->type, curr->type)) {
       replaceCurrent(curr->ref);
 
       // We must refinalize here, as we may be returning a more specific
@@ -1987,6 +1987,24 @@ struct OptimizeInstructions
       return;
     }
 
+    // The cast will not definitely succeed, but perhaps the heap type part of
+    // the cast will, at least. That would leave only nullability as an issue,
+    // that is, this means that the input ref is nullable but we are casting to
+    // non-null.
+    if (HeapType::isSubType(curr->ref->type.getHeapType(), intendedType)) {
+      assert(curr->ref->type.isNullable());
+      assert(curr->type.isNonNullable());
+
+      // Given the heap type will cast ok, all we need to do is check for a null
+      // here.
+      //
+      // As above, we must refinalize as we may now be emitting a more refined
+      // type (specifically a more refined heap type).
+      replaceCurrent(Builder(*getModule()).makeRefAs(RefAsNonNull, curr->ref));
+      refinalize = true;
+      return;
+    }
+
     // Repeated identical ref.cast operations are unnecessary. First, find the
     // immediate child cast, if there is one.
     // TODO: Look even further through incompatible casts?
diff --git a/test/lit/passes/optimize-instructions-gc.wast b/test/lit/passes/optimize-instructions-gc.wast
index 9ca0c81e0..04ae7a5cc 100644
--- a/test/lit/passes/optimize-instructions-gc.wast
+++ b/test/lit/passes/optimize-instructions-gc.wast
@@ -3120,4 +3120,67 @@
       (ref.null nofunc)
     )
   )
+
+  ;; CHECK:      (func $ref-cast-heap-type (type $ref?|$B|_ref|$B|_=>_none) (param $null-b (ref null $B)) (param $b (ref $B))
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $b)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $null-b)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $b)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.get $null-b)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $ref-cast-heap-type (type $ref?|$B|_ref|$B|_=>_none) (param $null-b (ref null $B)) (param $b (ref $B))
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (local.get $b)
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (ref.as_non_null
+  ;; NOMNL-NEXT:    (local.get $null-b)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (local.get $b)
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (local.get $null-b)
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  (func $ref-cast-heap-type (param $null-b (ref null $B)) (param $b (ref $B))
+    ;; We are casting a heap type to a supertype, which always succeeds. However
+    ;; we need to check for nullability.
+
+    ;; Non-nullable casts. When the input is non-nullable we must succeed.
+    (drop
+      (ref.cast $A
+        (local.get $b)
+      )
+    )
+    ;; When the input can be null, we might fail if it is a null. But we can
+    ;; switch to checking only that.
+    (drop
+      (ref.cast $A
+        (local.get $null-b)
+      )
+    )
+
+    ;; Null casts. Both of these must succeed.
+    (drop
+      (ref.cast null $A
+        (local.get $b)
+      )
+    )
+    (drop
+      (ref.cast null $A
+        (local.get $null-b)
+      )
+    )
+  )
 )