1 files changed, 14 insertions, 8 deletions
diff --git a/src/passes/CodeFolding.cpp b/src/passes/CodeFolding.cpp
index 21527da6b..0cddec4ca 100644
--- a/src/passes/CodeFolding.cpp
+++ b/src/passes/CodeFolding.cpp
@@ -84,7 +84,9 @@ struct ExpressionMarker
   void visitExpression(Expression* expr) { marked.insert(expr); }
 };
 
-struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
+struct CodeFolding
+  : public WalkerPass<
+      ControlFlowWalker<CodeFolding, UnifiedExpressionVisitor<CodeFolding>>> {
   bool isFunctionParallel() override { return true; }
 
   std::unique_ptr<Pass> create() override {
@@ -138,6 +140,17 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
 
   // walking
 
+  void visitExpression(Expression* curr) {
+    // For any branching instruction not explicitly handled by this pass, mark
+    // the labels it branches to unoptimizable.
+    // TODO: Handle folding br_on* instructions. br_on_null could be folded with
+    // other kinds of branches and br_on_non_null, br_on_cast, and
+    // br_on_cast_fail instructions could be folded with other copies of
+    // themselves.
+    BranchUtils::operateOnScopeNameUses(
+      curr, [&](Name label) { unoptimizables.insert(label); });
+  }
+
   void visitBreak(Break* curr) {
     if (curr->condition || curr->value) {
       unoptimizables.insert(curr->name);
@@ -155,13 +168,6 @@ struct CodeFolding : public WalkerPass<ControlFlowWalker<CodeFolding>> {
     }
   }
 
-  void visitSwitch(Switch* curr) {
-    for (auto target : curr->targets) {
-      unoptimizables.insert(target);
-    }
-    unoptimizables.insert(curr->default_);
-  }
-
   void visitUnreachable(Unreachable* curr) {
     // we can only optimize if we are at the end of the parent block
     if (!controlFlowStack.empty()) {