[Wasm GC] Optimize ref.as_non_null of a nullable cast (#5415)

2 files changed, 20 insertions, 4 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index 3265a1240..6ca314b96 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -2219,6 +2219,22 @@ struct OptimizeInstructions
 
     if (curr->op == RefAsNonNull && !curr->value->type.isNullable()) {
       replaceCurrent(curr->value);
+      return;
+    }
+
+    // As we do in visitRefCast, ref.cast can be combined with ref.as_non_null.
+    // This code handles the case where the ref.as is on the outside:
+    //
+    //   (ref.as_non_null (ref.cast null ..))
+    // =>
+    //   (ref.cast ..)
+    //
+    if (auto* cast = curr->value->dynCast<RefCast>()) {
+      // The cast cannot be non-nullable, or we would have handled this right
+      // above by just removing the ref.as, since it would not be needed.
+      assert(!cast->type.isNonNullable());
+      cast->type = Type(cast->type.getHeapType(), NonNullable);
+      replaceCurrent(cast);
     }
   }
 
diff --git a/test/lit/passes/optimize-instructions-gc-tnh.wast b/test/lit/passes/optimize-instructions-gc-tnh.wast
index 67f15de6c..948e03296 100644
--- a/test/lit/passes/optimize-instructions-gc-tnh.wast
+++ b/test/lit/passes/optimize-instructions-gc-tnh.wast
@@ -15,10 +15,8 @@
   ;; TNH-NEXT: )
   ;; NO_TNH:      (func $ref.eq (type $eqref_eqref_=>_i32) (param $a eqref) (param $b eqref) (result i32)
   ;; NO_TNH-NEXT:  (ref.eq
-  ;; NO_TNH-NEXT:   (ref.as_non_null
-  ;; NO_TNH-NEXT:    (ref.cast null $struct
-  ;; NO_TNH-NEXT:     (local.get $a)
-  ;; NO_TNH-NEXT:    )
+  ;; NO_TNH-NEXT:   (ref.cast $struct
+  ;; NO_TNH-NEXT:    (local.get $a)
   ;; NO_TNH-NEXT:   )
   ;; NO_TNH-NEXT:   (ref.as_data
   ;; NO_TNH-NEXT:    (local.get $b)
@@ -29,6 +27,8 @@
     ;; When traps never happen we can remove all the casts here, since they do
     ;; not affect the comparison of the references.
     (ref.eq
+      ;; When traps can happen we can still improve this by removing and
+      ;; combining redundant casts.
       (ref.as_data
         (ref.as_non_null
           (ref.cast null $struct