2 files changed, 120 insertions, 8 deletions

diff --git a/src/passes/MergeBlocks.cpp b/src/passes/MergeBlocks.cpp
index 4798ad28d..bde9397a8 100644
--- a/src/passes/MergeBlocks.cpp
+++ b/src/passes/MergeBlocks.cpp
@@ -64,9 +64,40 @@
 #include <wasm.h>
 #include <pass.h>
 #include <ast_utils.h>
+#include <wasm-builder.h>
 
 namespace wasm {
 
+struct SwitchFinder : public ControlFlowWalker<SwitchFinder, Visitor<SwitchFinder>> {
+  Expression* origin;
+  bool found = false;
+
+  void visitSwitch(Switch* curr) {
+    if (findBreakTarget(curr->default_) == origin) {
+      found = true;
+      return;
+    }
+    for (auto& target : curr->targets) {
+      if (findBreakTarget(target) == origin) {
+        found = true;
+        return;
+      }
+    }
+  }
+};
+
+struct BreakValueDropper : public ControlFlowWalker<BreakValueDropper, Visitor<BreakValueDropper>> {
+  Expression* origin;
+
+  void visitBreak(Break* curr) {
+    if (curr->value && findBreakTarget(curr->name) == origin) {
+      Builder builder(*getModule());
+      replaceCurrent(builder.makeSequence(builder.makeDrop(curr->value), curr));
+      curr->value = nullptr;
+    }
+  }
+};
+
 struct MergeBlocks : public WalkerPass<PostWalker<MergeBlocks, Visitor<MergeBlocks>>> {
   bool isFunctionParallel() override { return true; }
 
@@ -80,20 +111,39 @@ struct MergeBlocks : public WalkerPass<PostWalker<MergeBlocks, Visitor<MergeBloc
       for (size_t i = 0; i < curr->list.size(); i++) {
         Block* child = curr->list[i]->dynCast<Block>();
         if (!child) {
-          // TODO: if we have a child that is (drop (block ..)) then we can move the drop into the block, allowing more merging,
-          // but we must also drop values from brs
-          /*
+          // if we have a child that is (drop (block ..)) then we can move the drop into the block, and remove br values. this allows more merging,
           auto* drop = curr->list[i]->dynCast<Drop>();
           if (drop) {
             child = drop->value->dynCast<Block>();
             if (child) {
-              // reuse the drop
-              drop->value = child->list.back();
-              child->list.back() = drop;
-              curr->list[i] = child;
+              if (child->name.is()) {
+                Expression* expression = child;
+                // if there is a switch targeting us, we can't do it - we can't remove the value from other targets too
+                SwitchFinder finder;
+                finder.origin = child;
+                finder.walk(expression);
+                if (finder.found) {
+                  child = nullptr;
+                } else {
+                  // fix up breaks
+                  BreakValueDropper fixer;
+                  fixer.origin = child;
+                  fixer.setModule(getModule());
+                  fixer.walk(expression);
+                }
+              }
+              if (child) {
+                // we can do it!
+                // reuse the drop
+                drop->value = child->list.back();
+                child->list.back() = drop;
+                child->finalize();
+                curr->list[i] = child;
+                more = true;
+                changed = true;
+              }
             }
           }
-          */
         }
         if (!child) continue;
         if (child->name.is()) continue; // named blocks can have breaks to them (and certainly do, if we ran RemoveUnusedNames and RemoveUnusedBrs)
diff --git a/src/wasm-traversal.h b/src/wasm-traversal.h
index f9d3d8413..f4483644c 100644
--- a/src/wasm-traversal.h
+++ b/src/wasm-traversal.h
@@ -446,6 +446,68 @@ struct PostWalker : public Walker<SubType, VisitorType> {
   }
 };
 
+// Traversal with a control-flow stack.
+
+template<typename SubType, typename VisitorType>
+struct ControlFlowWalker : public PostWalker<SubType, VisitorType> {
+  ControlFlowWalker() {}
+
+  std::vector<Expression*> controlFlowStack; // contains blocks, loops, and ifs
+
+  // Uses the control flow stack to find the target of a break to a name
+  Expression* findBreakTarget(Name name) {
+    assert(!controlFlowStack.empty());
+    Index i = controlFlowStack.size() - 1;
+    while (1) {
+      auto* curr = controlFlowStack[i];
+      if (Block* block = curr->dynCast<Block>()) {
+        if (name == block->name) return curr;
+      } else if (Loop* loop = curr->dynCast<Loop>()) {
+        if (name == loop->name) return curr;
+      } else {
+        WASM_UNREACHABLE();
+      }
+      if (i == 0) return nullptr;
+      i--;
+    }
+  }
+
+  static void doPreVisitControlFlow(SubType* self, Expression** currp) {
+    self->controlFlowStack.push_back(*currp);
+  }
+
+  static void doPostVisitControlFlow(SubType* self, Expression** currp) {
+    assert(self->controlFlowStack.back() == *currp);
+    self->controlFlowStack.pop_back();
+  }
+
+  static void scan(SubType* self, Expression** currp) {
+    auto* curr = *currp;
+
+    switch (curr->_id) {
+      case Expression::Id::BlockId:
+      case Expression::Id::IfId:
+      case Expression::Id::LoopId: {
+        self->pushTask(SubType::doPostVisitControlFlow, currp);
+        break;
+      }
+      default: {}
+    }
+
+    PostWalker<SubType, VisitorType>::scan(self, currp);
+
+    switch (curr->_id) {
+      case Expression::Id::BlockId:
+      case Expression::Id::IfId:
+      case Expression::Id::LoopId: {
+        self->pushTask(SubType::doPreVisitControlFlow, currp);
+        break;
+      }
+      default: {}
+    }
+  }
+};
+
 // Traversal in the order of execution. This is quick and simple, but
 // does not provide the same comprehensive information that a full
 // conversion to basic blocks would. What it does give is a quick