[WasmGC] Fix try-catch and non-nullable local sets in values in SimplifyLocals (#4705)

Followup to #4703, this also handles the case where there is a non-
nullable local.set in the value of a nullable one, which we also cannot
optimize.

Fixes #4702

2 files changed, 85 insertions, 6 deletions

diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index 5cf3b8241..159f09963 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -306,8 +306,12 @@ struct SimplifyLocals
       for (auto& [index, info] : self->sinkables) {
         // 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
+        if (info.effects.throws()) {
+          invalidated.push_back(index);
+          continue;
+        }
+
+        // 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)
@@ -338,9 +342,17 @@ struct SimplifyLocals
         // 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);
+        //
+        // The problem described can also occur on the *value* and not the set
+        // itself. For example, |X| above could be a local.set of a non-nullable
+        // local. For that reason we must scan it all.
+        if (self->getModule()->features.hasGCNNLocals()) {
+          for (auto* set : FindAll<LocalSet>(*info.item).list) {
+            if (self->getFunction()->getLocalType(set->index).isNonNullable()) {
+              invalidated.push_back(index);
+              break;
+            }
+          }
         }
       }
       for (auto index : invalidated) {
diff --git a/test/lit/passes/simplify-locals-gc-nn.wast b/test/lit/passes/simplify-locals-gc-nn.wast
index eec97a4d4..9d75b3d80 100644
--- a/test/lit/passes/simplify-locals-gc-nn.wast
+++ b/test/lit/passes/simplify-locals-gc-nn.wast
@@ -87,4 +87,71 @@
         )
       )
     )
-  ))
+  )
+
+  ;; CHECK:      (func $test-nested-nn
+  ;; CHECK-NEXT:  (local $nullable anyref)
+  ;; CHECK-NEXT:  (local $nn (ref any))
+  ;; CHECK-NEXT:  (local.set $nullable
+  ;; CHECK-NEXT:   (block $block (result (ref any))
+  ;; CHECK-NEXT:    (local.set $nn
+  ;; CHECK-NEXT:     (ref.as_non_null
+  ;; CHECK-NEXT:      (ref.null any)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.as_non_null
+  ;; CHECK-NEXT:     (ref.null any)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (try $try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $nullable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch_all
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $nullable)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (local.get $nn)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test-nested-nn
+    ;; As above, the set we want to move is nullable, but now it has a nested
+    ;; value that is a non-nullable set. That should also prevent us from
+    ;; moving it.
+    (local $nullable (ref null any))
+    (local $nn (ref any))
+    (local.set $nullable
+      (block (result (ref any))
+        (local.set $nn
+          (ref.as_non_null
+            (ref.null any)
+          )
+        )
+        (ref.as_non_null
+          (ref.null any)
+        )
+      )
+    )
+    (try
+      (do
+        (drop
+          (local.get $nullable)
+        )
+      )
+      (catch_all
+        (drop
+          (local.get $nullable)
+        )
+        (drop
+          (local.get $nn)
+        )
+      )
+    )
+  )
+)