1 files changed, 23 insertions, 6 deletions
diff --git a/src/wasm-traversal.h b/src/wasm-traversal.h
index 522f5092c..f49407cad 100644
--- a/src/wasm-traversal.h
+++ b/src/wasm-traversal.h
@@ -157,7 +157,11 @@ struct Walker : public VisitorType {
   // just one visit*() method is called by the traversal; if you replace a node,
   // and you want to process the output, you must do that explicitly).
   Expression* replaceCurrent(Expression* expression) {
-    return replace = expression;
+    return *replacep = expression;
+  }
+
+  Expression* getCurrent() {
+    return *replacep;
   }
 
   // Get the current module
@@ -184,6 +188,15 @@ struct Walker : public VisitorType {
     setFunction(nullptr);
   }
 
+  void walkFunctionInModule(Function* func, Module* module) {
+    setModule(module);
+    setFunction(func);
+    static_cast<SubType*>(this)->doWalkFunction(func);
+    static_cast<SubType*>(this)->visitFunction(func);
+    setFunction(nullptr);
+    setModule(nullptr);
+  }
+
   // override this to provide custom functionality
   void doWalkFunction(Function* func) {
     walk(func->body);
@@ -264,12 +277,9 @@ struct Walker : public VisitorType {
     pushTask(SubType::scan, &root);
     while (stack.size() > 0) {
       auto task = popTask();
+      replacep = task.currp;
       assert(*task.currp);
       task.func(static_cast<SubType*>(this), task.currp);
-      if (replace) {
-        *task.currp = replace;
-        replace = nullptr;
-      }
     }
   }
 
@@ -311,7 +321,7 @@ struct Walker : public VisitorType {
   }
 
 private:
-  Expression* replace = nullptr; // a node to replace
+  Expression** replacep = nullptr; // the address of the current node, used to replace it
   std::vector<Task> stack; // stack of tasks
   Function* currFunction = nullptr; // current function being processed
   Module* currModule = nullptr; // current module being processed
@@ -571,6 +581,13 @@ struct ExpressionStackWalker : public PostWalker<SubType, VisitorType> {
 
     self->pushTask(SubType::doPreVisit, currp);
   }
+
+  Expression* replaceCurrent(Expression* expression) {
+    PostWalker<SubType, VisitorType>::replaceCurrent(expression);
+    // also update the stack
+    expressionStack.back() = expression;
+    return expression;
+  }
 };
 
 // Traversal in the order of execution. This is quick and simple, but