2 files changed, 56 insertions, 40 deletions

diff --git a/src/cfg/cfg-traversal.h b/src/cfg/cfg-traversal.h
index 3f2eecb88..f66dc8d3f 100644
--- a/src/cfg/cfg-traversal.h
+++ b/src/cfg/cfg-traversal.h
@@ -30,6 +30,7 @@
 #ifndef cfg_traversal_h
 #define cfg_traversal_h
 
+#include "ir/branch-utils.h"
 #include "wasm-traversal.h"
 #include "wasm.h"
 
@@ -203,11 +204,15 @@ struct CFGWalker : public ControlFlowWalker<SubType, VisitorType> {
     self->loopStack.pop_back();
   }
 
-  static void doEndBreak(SubType* self, Expression** currp) {
-    auto* curr = (*currp)->cast<Break>();
-    self->branches[self->findBreakTarget(curr->name)].push_back(
-      self->currBasicBlock); // branch to the target
-    if (curr->condition) {
+  static void doEndBranch(SubType* self, Expression** currp) {
+    auto* curr = *currp;
+    auto branchTargets = BranchUtils::getUniqueTargets(curr);
+    // Add branches to the targets.
+    for (auto target : branchTargets) {
+      self->branches[self->findBreakTarget(target)].push_back(
+        self->currBasicBlock);
+    }
+    if (curr->type != Type::unreachable) {
       auto* last = self->currBasicBlock;
       self->link(last, self->startBasicBlock()); // we might fall through
     } else {
@@ -215,24 +220,6 @@ struct CFGWalker : public ControlFlowWalker<SubType, VisitorType> {
     }
   }
 
-  static void doEndSwitch(SubType* self, Expression** currp) {
-    auto* curr = (*currp)->cast<Switch>();
-    // we might see the same label more than once; do not spam branches
-    std::set<Name> seen;
-    for (Name target : curr->targets) {
-      if (!seen.count(target)) {
-        self->branches[self->findBreakTarget(target)].push_back(
-          self->currBasicBlock); // branch to the target
-        seen.insert(target);
-      }
-    }
-    if (!seen.count(curr->default_)) {
-      self->branches[self->findBreakTarget(curr->default_)].push_back(
-        self->currBasicBlock); // branch to the target
-    }
-    self->startUnreachableBlock();
-  }
-
   static void doEndThrowingInst(SubType* self, Expression** currp) {
     // Even if the instruction can possibly throw, we don't end the current
     // basic block unless the instruction is within a try-catch, because the CFG
@@ -366,22 +353,6 @@ struct CFGWalker : public ControlFlowWalker<SubType, VisitorType> {
         self->pushTask(SubType::doEndLoop, currp);
         break;
       }
-      case Expression::Id::BreakId: {
-        self->pushTask(SubType::doEndBreak, currp);
-        break;
-      }
-      case Expression::Id::SwitchId: {
-        self->pushTask(SubType::doEndSwitch, currp);
-        break;
-      }
-      case Expression::Id::ReturnId: {
-        self->pushTask(SubType::doStartUnreachableBlock, currp);
-        break;
-      }
-      case Expression::Id::UnreachableId: {
-        self->pushTask(SubType::doStartUnreachableBlock, currp);
-        break;
-      }
       case Expression::Id::CallId:
       case Expression::Id::CallIndirectId: {
         self->pushTask(SubType::doEndCall, currp);
@@ -405,7 +376,13 @@ struct CFGWalker : public ControlFlowWalker<SubType, VisitorType> {
         self->pushTask(SubType::doEndThrow, currp);
         break;
       }
-      default: {}
+      default: {
+        if (Properties::isBranch(curr)) {
+          self->pushTask(SubType::doEndBranch, currp);
+        } else if (curr->type == Type::unreachable) {
+          self->pushTask(SubType::doStartUnreachableBlock, currp);
+        }
+      }
     }
 
     ControlFlowWalker<SubType, VisitorType>::scan(self, currp);
diff --git a/test/lit/passes/coalesce-locals-gc.wast b/test/lit/passes/coalesce-locals-gc.wast
index 5bad35001..173ad02be 100644
--- a/test/lit/passes/coalesce-locals-gc.wast
+++ b/test/lit/passes/coalesce-locals-gc.wast
@@ -3,6 +3,9 @@
 ;; RUN:   | filecheck %s
 
 (module
+ (type $array (array (mut i8)))
+ (global $global (ref null $array) (ref.null $array))
+
  ;; CHECK:      (func $test-dead-get-non-nullable (param $0 dataref)
  ;; CHECK-NEXT:  (unreachable)
  ;; CHECK-NEXT:  (drop
@@ -18,4 +21,40 @@
    (local.get $func)
   )
  )
+
+ ;; CHECK:      (func $br_on_null (param $0 (ref null $array)) (result (ref null $array))
+ ;; CHECK-NEXT:  (block $label$1 (result (ref null $array))
+ ;; CHECK-NEXT:   (block $label$2
+ ;; CHECK-NEXT:    (br $label$1
+ ;; CHECK-NEXT:     (br_on_null $label$2
+ ;; CHECK-NEXT:      (local.get $0)
+ ;; CHECK-NEXT:     )
+ ;; CHECK-NEXT:    )
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (local.set $0
+ ;; CHECK-NEXT:    (global.get $global)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (local.get $0)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $br_on_null (param $ref (ref null $array)) (result (ref null $array))
+  (local $1 (ref null $array))
+  (block $label$1 (result (ref null $array))
+   (block $label$2
+    (br $label$1
+     ;; Test that we properly model the basic block connections around a
+     ;; BrOnNull. There should be a branch to $label$2, and also a fallthrough.
+     ;; As a result, the local.set below is reachable, and should not be
+     ;; eliminated (turned into a drop).
+     (br_on_null $label$2
+      (local.get $ref)
+     )
+    )
+   )
+   (local.set $1
+    (global.get $global)
+   )
+   (local.get $1)
+  )
+ )
 )