[Wasm GC] Fix LinearExecutionWalker (#3954)

That traversal did not mention BrOn, which led to it doing incorrect work in
SimplifyLocals.

Also add assertions at the end, that aim to prevent future issues.

The rest of the fix is to make SimplifyLocals not assume that things
are a Switch if they are not an If/Block/etc., so that we don't crash
on a BrOn.

3 files changed, 71 insertions, 3 deletions

diff --git a/src/ir/linear-execution.h b/src/ir/linear-execution.h
index b1d0b96c3..a37349489 100644
--- a/src/ir/linear-execution.h
+++ b/src/ir/linear-execution.h
@@ -17,6 +17,7 @@
 #ifndef wasm_ir_linear_execution_h
 #define wasm_ir_linear_execution_h
 
+#include <ir/properties.h>
 #include <wasm-traversal.h>
 #include <wasm.h>
 
@@ -126,7 +127,17 @@ struct LinearExecutionWalker : public PostWalker<SubType, VisitorType> {
         self->pushTask(SubType::doNoteNonLinear, currp);
         break;
       }
+      case Expression::Id::BrOnId: {
+        self->pushTask(SubType::doVisitBrOn, currp);
+        self->pushTask(SubType::doNoteNonLinear, currp);
+        self->maybePushTask(SubType::scan, &curr->cast<BrOn>()->rtt);
+        self->pushTask(SubType::scan, &curr->cast<BrOn>()->ref);
+        break;
+      }
       default: {
+        // All relevant things should have been handled.
+        assert(!Properties::isControlFlowStructure(curr));
+        assert(!Properties::isBranch(curr));
         // other node types do not have control flow, use regular post-order
         PostWalker<SubType, VisitorType>::scan(self, currp);
       }
diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index 02be1ceec..77b8f1a93 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -140,9 +140,11 @@ struct SimplifyLocals
     } else if (curr->is<If>()) {
       assert(!curr->cast<If>()
                 ->ifFalse); // if-elses are handled by doNoteIf* methods
-    } else if (curr->is<Switch>()) {
-      auto* sw = curr->cast<Switch>();
-      auto targets = BranchUtils::getUniqueTargets(sw);
+    } else {
+      // Not one of the recognized instructions, so do not optimize here: mark
+      // all the targets as unoptimizable.
+      // TODO optimize BrOn, Switch, etc.
+      auto targets = BranchUtils::getUniqueTargets(curr);
       for (auto target : targets) {
         self->unoptimizableBlocks.insert(target);
       }
diff --git a/test/lit/passes/simplify-locals-gc.wast b/test/lit/passes/simplify-locals-gc.wast
index 06822a5d7..89e696d8f 100644
--- a/test/lit/passes/simplify-locals-gc.wast
+++ b/test/lit/passes/simplify-locals-gc.wast
@@ -33,4 +33,59 @@
     )
     (local.get $temp)
   )
+
+  ;; CHECK:      (func $no-block-values-if-br_on
+  ;; CHECK-NEXT:  (local $temp anyref)
+  ;; CHECK-NEXT:  (block $block
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (br_on_null $block
+  ;; CHECK-NEXT:     (ref.null any)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $temp
+  ;; CHECK-NEXT:    (ref.null any)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (br $block)
+  ;; CHECK-NEXT:   (local.set $temp
+  ;; CHECK-NEXT:    (ref.null any)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.as_non_null
+  ;; CHECK-NEXT:    (local.get $temp)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $no-block-values-if-br_on
+   (local $temp (ref null any))
+   (block $block
+    (drop
+     ;; This br_on should inhibit trying to create a block return value for
+     ;; this block. Aside from the br_on, it looks correct, i.e., we have a
+     ;; break with a set before it, and a set before the end of the block. Due
+     ;; to the br_on's presence, the pass should not do anything to this
+     ;; function.
+     ;;
+     ;; TODO: support br_on in this optimization eventually, but the variable
+     ;;       possible return values and sent values make that nontrivial.
+     (br_on_null $block
+      (ref.null any)
+     )
+    )
+    (local.set $temp
+     (ref.null any)
+    )
+    (br $block)
+    (local.set $temp
+     (ref.null any)
+    )
+   )
+   ;; Attempt to use the local that the pass will try to move to a block return
+   ;; value, to cause the optimization to try to run.
+   (drop
+    (ref.as_non_null
+     (local.get $temp)
+    )
+   )
+  )
 )