[Wasm GC] Optimize away ref.as_non_null going into local.set in TNH mode (#4157)

If we can remove such traps, we can remove ref.as_non_null if the local
type is nullable anyhow.

If we support non-nullable locals, however, then do not do this, as it
could inhibit specializing the local type later. Do the same for tees which
we had existing code for.

Background: #4061 (comment)

4 files changed, 89 insertions, 11 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index c9e4d7ff5..b9b426466 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -1054,21 +1054,36 @@ struct OptimizeInstructions
   }
 
   void visitLocalSet(LocalSet* curr) {
-    //   (local.tee (ref.as_non_null ..))
-    // can be reordered to
-    //   (ref.as_non_null (local.tee ..))
-    // if the local is nullable (which it must be until some form of let is
-    // added). The reordering allows the ref.as to be potentially optimized
-    // further based on where the value flows to.
-    if (curr->isTee()) {
-      if (auto* as = curr->value->dynCast<RefAs>()) {
-        if (as->op == RefAsNonNull &&
-            getFunction()->getLocalType(curr->index).isNullable()) {
+    // Interactions between local.set/tee and ref.as_non_null can be optimized
+    // in some cases, by removing or moving the ref.as_non_null operation. In
+    // all cases, we only do this when we do *not* allow non-nullable locals. If
+    // we do allow such locals, then (1) this local might be non-nullable, so we
+    // can't remove or move a ref.as_non_null flowing into a local.set/tee, and
+    // (2) even if the local were nullable, if we change things we might prevent
+    // the LocalSubtyping pass from turning it into a non-nullable local later.
+    if (auto* as = curr->value->dynCast<RefAs>()) {
+      if (as->op == RefAsNonNull && !getModule()->features.hasGCNNLocals()) {
+        //   (local.tee (ref.as_non_null ..))
+        // =>
+        //   (ref.as_non_null (local.tee ..))
+        //
+        // The reordering allows the ref.as to be potentially optimized further
+        // based on where the value flows to.
+        if (curr->isTee()) {
           curr->value = as->value;
           curr->finalize();
           as->value = curr;
           as->finalize();
           replaceCurrent(as);
+          return;
+        }
+
+        // Otherwise, if this is not a tee, then no value falls through. The
+        // ref.as_non_null acts as a null check here, basically. If we are
+        // ignoring such traps, we can remove it.
+        auto passOptions = getPassOptions();
+        if (passOptions.ignoreImplicitTraps || passOptions.trapsNeverHappen) {
+          curr->value = as->value;
         }
       }
     }
diff --git a/test/lit/passes/optimize-instructions-gc-iit.wast b/test/lit/passes/optimize-instructions-gc-iit.wast
index 41d4e69d5..b99042643 100644
--- a/test/lit/passes/optimize-instructions-gc-iit.wast
+++ b/test/lit/passes/optimize-instructions-gc-iit.wast
@@ -342,4 +342,28 @@
       )
     )
   )
+
+  ;; CHECK:      (func $set-of-as-non-null (param $x anyref)
+  ;; CHECK-NEXT:  (local.set $x
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $set-of-as-non-null (param $x anyref)
+  ;; NOMNL-NEXT:  (local.set $x
+  ;; NOMNL-NEXT:   (local.get $x)
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  ;; NOMNL-TNH:      (func $set-of-as-non-null (param $x anyref)
+  ;; NOMNL-TNH-NEXT:  (local.set $x
+  ;; NOMNL-TNH-NEXT:   (local.get $x)
+  ;; NOMNL-TNH-NEXT:  )
+  ;; NOMNL-TNH-NEXT: )
+  (func $set-of-as-non-null (param $x anyref)
+    ;; As we ignore such traps, we can remove the ref.as here.
+    (local.set $x
+      (ref.as_non_null
+        (local.get $x)
+      )
+    )
+  )
 )
diff --git a/test/lit/passes/optimize-instructions-gc-tnh-nn.wast b/test/lit/passes/optimize-instructions-gc-tnh-nn.wast
new file mode 100644
index 000000000..e0a580ae1
--- /dev/null
+++ b/test/lit/passes/optimize-instructions-gc-tnh-nn.wast
@@ -0,0 +1,40 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s --optimize-instructions --traps-never-happen --enable-gc-nn-locals -all -S -o - \
+;; RUN:   | filecheck %s
+
+(module
+  ;; CHECK:      (func $set-of-as-non-null
+  ;; CHECK-NEXT:  (local $x anyref)
+  ;; CHECK-NEXT:  (local.set $x
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $x)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (local.tee $x
+  ;; CHECK-NEXT:    (ref.as_non_null
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $set-of-as-non-null
+    (local $x anyref)
+    ;; As we ignore such traps, we can in principle remove the ref.as here.
+    ;; However, as we allow non-nullable locals, we should not do that - if we
+    ;; did it it might prevent specializing the local type later.
+    (local.set $x
+      (ref.as_non_null
+        (local.get $x)
+      )
+    )
+    ;; The same for a tee.
+    (drop
+      (local.tee $x
+        (ref.as_non_null
+          (local.get $x)
+        )
+      )
+    )
+  )
+)
diff --git a/test/lit/passes/optimize-instructions-gc.wast b/test/lit/passes/optimize-instructions-gc.wast
index a3f91c756..3f12f58e9 100644
--- a/test/lit/passes/optimize-instructions-gc.wast
+++ b/test/lit/passes/optimize-instructions-gc.wast
@@ -1247,7 +1247,6 @@
         (ref.as_non_null
           (local.get $x)
         )
-        (rtt.canon $struct)
       )
     )
   )