[WasmGC] Fix non-nullable locals and try-catch in SimplifyLocals (#4703)

Binaryen will not change dominance in SimplifyLocals, however, the current spec's
notion of dominance is simpler than ours, and we must not optimize certain cases in
order to still validate. See details in the comment and test.

Helps #4702

2 files changed, 128 insertions, 3 deletions

diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index 35980f503..5cf3b8241 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -299,12 +299,47 @@ struct SimplifyLocals
            Expression** currp) {
     Expression* curr = *currp;
 
-    // Expressions that may throw cannot be sinked into 'try'. At the start of
-    // 'try', we drop all sinkables that may throw.
+    // Certain expressions cannot be sinked into 'try', and so at the start of
+    // 'try' we forget about them.
     if (curr->is<Try>()) {
       std::vector<Index> invalidated;
       for (auto& [index, info] : self->sinkables) {
-        if (info.effects.throws()) {
+        // Expressions that may throw cannot be moved into a try (which might
+        // catch them, unlike before the move).
+        //
+        // Also, non-nullable local.sets cannot be moved into a try, as that may
+        // change dominance from the perspective of the spec
+        //
+        //  (local.set $x X)
+        //  (try
+        //    ..
+        //    (Y
+        //      (local.get $x))
+        //  (catch
+        //    (Z
+        //      (local.get $x)))
+        //
+        // =>
+        //
+        //  (try
+        //    ..
+        //    (Y
+        //      (local.tee $x X))
+        //  (catch
+        //    (Z
+        //      (local.get $x)))
+        //
+        // After sinking the set, the tee does not dominate the get in the
+        // catch, at least not in the simple way the spec defines it, see
+        // https://github.com/WebAssembly/function-references/issues/44#issuecomment-1083146887
+        // We have more refined information about control flow and dominance
+        // than the spec, and so we would see if ".." can throw or not (only if
+        // it can throw is there a branch to the catch, which can change
+        // dominance). To stay compliant with the spec, however, we must not
+        // move code regardless of whether ".." can throw - we must simply keep
+        // the set outside of the try.
+        if (info.effects.throws() ||
+            self->getFunction()->getLocalType(index).isNonNullable()) {
           invalidated.push_back(index);
         }
       }
diff --git a/test/lit/passes/simplify-locals-gc-nn.wast b/test/lit/passes/simplify-locals-gc-nn.wast
new file mode 100644
index 000000000..eec97a4d4
--- /dev/null
+++ b/test/lit/passes/simplify-locals-gc-nn.wast
@@ -0,0 +1,90 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s --simplify-locals -all --enable-gc-nn-locals -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (func $test-nn
+  ;; CHECK-NEXT:  (local $nn (ref any))
+  ;; CHECK-NEXT:  (local.set $nn
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (ref.null any)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (try $try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $nn)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch_all
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $nn)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test-nn
+    (local $nn (ref any))
+    ;; We should not sink this set into the try, as the spec does not allow it
+    ;; even though we are not changing dominance (we are not changing it,
+    ;; because there is nothing that can throw in the try's body that can reach
+    ;; the catch_all before the local.set that we move there). See
+    ;; https://github.com/WebAssembly/function-references/issues/44#issuecomment-1083146887
+    (local.set $nn
+      (ref.as_non_null
+        (ref.null any)
+      )
+    )
+    (try
+      (do
+        (drop
+          (local.get $nn)
+        )
+      )
+      (catch_all
+        (drop
+          (local.get $nn)
+        )
+      )
+    )
+  )
+
+  ;; CHECK:      (func $test-nullable
+  ;; CHECK-NEXT:  (local $nullable anyref)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (try $try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (local.set $nullable
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (ref.null any)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch_all
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $nullable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test-nullable
+    ;; As above, but now the local is nullable. Here we *can* optimize the set
+    ;; into the try.
+    (local $nullable (ref null any))
+    (local.set $nullable
+      (ref.as_non_null
+        (ref.null any)
+      )
+    )
+    (try
+      (do
+        (drop
+          (local.get $nullable)
+        )
+      )
+      (catch_all
+        (drop
+          (local.get $nullable)
+        )
+      )
+    )
+  ))