2 files changed, 44 insertions, 20 deletions

diff --git a/src/passes/Vacuum.cpp b/src/passes/Vacuum.cpp
index 1b0887181..bb005e36c 100644
--- a/src/passes/Vacuum.cpp
+++ b/src/passes/Vacuum.cpp
@@ -25,13 +25,11 @@
 
 namespace wasm {
 
-struct Vacuum : public WalkerPass<PostWalker<Vacuum, Visitor<Vacuum>>> {
+struct Vacuum : public WalkerPass<ExpressionStackWalker<Vacuum, Visitor<Vacuum>>> {
   bool isFunctionParallel() override { return true; }
 
   Pass* create() override { return new Vacuum; }
 
-  std::vector<Expression*> expressionStack;
-
   // returns nullptr if curr is dead, curr if it must stay as is, or another node if it can be replaced
   Expression* optimize(Expression* curr, bool resultUsed) {
     while (1) {
@@ -206,23 +204,6 @@ struct Vacuum : public WalkerPass<PostWalker<Vacuum, Visitor<Vacuum>>> {
     }
   }
 
-  static void visitPre(Vacuum* self, Expression** currp) {
-    self->expressionStack.push_back(*currp);
-  }
-
-  static void visitPost(Vacuum* self, Expression** currp) {
-    self->expressionStack.pop_back();
-  }
-
-  // override scan to add a pre and a post check task to all nodes
-  static void scan(Vacuum* self, Expression** currp) {
-    self->pushTask(visitPost, currp);
-
-    WalkerPass<PostWalker<Vacuum, Visitor<Vacuum>>>::scan(self, currp);
-
-    self->pushTask(visitPre, currp);
-  }
-
   void visitFunction(Function* curr) {
     auto* optimized = optimize(curr->body, curr->result != none);
     if (optimized) {
diff --git a/src/wasm-traversal.h b/src/wasm-traversal.h
index b69db1ea2..a345a7f4b 100644
--- a/src/wasm-traversal.h
+++ b/src/wasm-traversal.h
@@ -508,6 +508,49 @@ struct ControlFlowWalker : public PostWalker<SubType, VisitorType> {
   }
 };
 
+// Traversal with an expression stack.
+
+template<typename SubType, typename VisitorType>
+struct ExpressionStackWalker : public PostWalker<SubType, VisitorType> {
+  ExpressionStackWalker() {}
+
+  std::vector<Expression*> expressionStack;
+
+  // Uses the control flow stack to find the target of a break to a name
+  Expression* findBreakTarget(Name name) {
+    assert(!expressionStack.empty());
+    Index i = expressionStack.size() - 1;
+    while (1) {
+      auto* curr = expressionStack[i];
+      if (Block* block = curr->template dynCast<Block>()) {
+        if (name == block->name) return curr;
+      } else if (Loop* loop = curr->template dynCast<Loop>()) {
+        if (name == loop->name) return curr;
+      } else {
+        WASM_UNREACHABLE();
+      }
+      if (i == 0) return nullptr;
+      i--;
+    }
+  }
+
+  static void doPreVisit(SubType* self, Expression** currp) {
+    self->expressionStack.push_back(*currp);
+  }
+
+  static void doPostVisit(SubType* self, Expression** currp) {
+    self->expressionStack.pop_back();
+  }
+
+  static void scan(SubType* self, Expression** currp) {
+    self->pushTask(SubType::doPostVisit, currp);
+
+    PostWalker<SubType, VisitorType>::scan(self, currp);
+
+    self->pushTask(SubType::doPreVisit, currp);
+  }
+};
+
 // 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