[WasmGC] OptimizeInstructions: Reorder externalize/internalize operations with ref.as_non_null (#7004)

 (any.convert_extern/extern.convert_any (ref.as_non_null ..))
 =>
 (ref.as_non_null (any.convert_extern/extern.convert_any ..))

This then allows the RefAsNonNull to be combined with parents in some cases
(whereas the reverse allows nothing).

Author: kripken

src/passes/OptimizeInstructions.cpp

@@ -2248,8 +2248,27 @@ struct OptimizeInstructions
     }
 
     if (curr->op == ExternConvertAny || curr->op == AnyConvertExtern) {
-      // We can't optimize these. Even removing a non-null cast is not valid as
-      // they allow nulls to filter through, unlike other RefAs*.
+      // These pass nulls through, and we can reorder them with null traps:
+      //
+      //  (any.convert_extern/extern.convert_any (ref.as_non_null.. ))
+      // =>
+      //  (ref.as_non_null (any.convert_extern/extern.convert_any ..))
+      //
+      // By moving the RefAsNonNull outside, it may reach a position where it
+      // can be optimized (e.g. if the parent traps anyhow). And,
+      // ExternConvertAny/AnyConvertExtern cannot be folded with anything, so
+      // there is no harm to moving them inside.
+      if (auto* refAsChild = curr->value->dynCast<RefAs>()) {
+        if (refAsChild->op == RefAsNonNull) {
+          // Reorder and fix up the types.
+          curr->value = refAsChild->value;
+          curr->finalize();
+          refAsChild->value = curr;
+          refAsChild->finalize();
+          replaceCurrent(refAsChild);
+        }
+      }
+      // TODO: optimize away ExternConvertAny of AnyConvertExtern, etc.
       return;
     }
 

test/lit/passes/optimize-instructions-gc-extern.wast

@@ -3,15 +3,18 @@
 ;; RUN:   | filecheck %s
 
 (module
+  ;; CHECK:      (type $array (array (mut i8)))
+  (type $array (array (mut i8)))
+
   ;; CHECK:      (func $extern.convert_any (type $0) (param $x anyref) (param $y externref)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (extern.convert_any
   ;; CHECK-NEXT:    (local.get $x)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (extern.convert_any
-  ;; CHECK-NEXT:    (ref.as_non_null
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (extern.convert_any
   ;; CHECK-NEXT:     (local.get $x)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
@@ -22,27 +25,30 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (any.convert_extern
-  ;; CHECK-NEXT:    (ref.as_non_null
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (any.convert_extern
   ;; CHECK-NEXT:     (local.get $y)
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
-  (func $extern.convert_any (export "ext") (param $x (ref null any)) (param $y (ref null extern))
+  (func $extern.convert_any (param $x (ref null any)) (param $y (ref null extern))
     ;; We should not change anything here, and also not hit an internal error.
     (drop
       (extern.convert_any
         (local.get $x)
       )
     )
+    ;; We can reorder the externalize with the ref.as_non_null, which sometimes
+    ;; helps later optimizations, see below.
     (drop
       (extern.convert_any
         (ref.as_non_null
           (local.get $x)
         )
       )
     )
+    ;; As the above two cases, but for internalize.
     (drop
       (any.convert_extern
         (local.get $y)
@@ -56,4 +62,27 @@
       )
     )
   )
+
+  ;; CHECK:      (func $convert.optimize.parent (type $1) (param $ext externref)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.cast (ref $array)
+  ;; CHECK-NEXT:    (any.convert_extern
+  ;; CHECK-NEXT:     (local.get $ext)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $convert.optimize.parent (param $ext externref)
+    ;; The ref.cast can fold in the ref.as_non_null, after it is moved
+    ;; outside of the any.convert_extern.
+    (drop
+      (ref.cast (ref null $array)
+        (any.convert_extern
+          (ref.as_non_null
+            (local.get $ext)
+          )
+        )
+      )
+    )
+  )
 )