refactor walk logic into walk* and doWalk* methods, for a more regular API that is clearer where it should be overridden (#551)

9 files changed, 41 insertions, 30 deletions

diff --git a/src/cfg/cfg-traversal.h b/src/cfg/cfg-traversal.h
index 74afc7525..d690de4aa 100644
--- a/src/cfg/cfg-traversal.h
+++ b/src/cfg/cfg-traversal.h
@@ -241,12 +241,12 @@ struct CFGWalker : public PostWalker<SubType, VisitorType> {
     }
   }
 
-  void walk(Expression*& root) {
+  void doWalkFunction(Function* func) {
     basicBlocks.clear();
 
     doStartBasicBlock(static_cast<SubType*>(this), nullptr);
     entry = currBasicBlock;
-    PostWalker<SubType, VisitorType>::walk(root);
+    PostWalker<SubType, VisitorType>::doWalkFunction(func);
 
     assert(branches.size() == 0);
     assert(ifStack.size() == 0);
diff --git a/src/pass.h b/src/pass.h
index bfbcb26d8..e25809faa 100644
--- a/src/pass.h
+++ b/src/pass.h
@@ -135,7 +135,7 @@ class WalkerPass : public Pass, public WalkerType {
 public:
   void run(PassRunner* runner, Module* module) override {
     prepare(runner, module);
-    WalkerType::startWalk(module);
+    WalkerType::walkModule(module);
     finalize(runner, module);
   }
 };
diff --git a/src/passes/CoalesceLocals.cpp b/src/passes/CoalesceLocals.cpp
index 17f6b08c3..39c480107 100644
--- a/src/passes/CoalesceLocals.cpp
+++ b/src/passes/CoalesceLocals.cpp
@@ -177,7 +177,7 @@ struct CoalesceLocals : public WalkerPass<CFGWalker<CoalesceLocals, Visitor<Coal
 
   // main entry point
 
-  void walk(Expression*& root);
+  void doWalkFunction(Function* func);
 
   void flowLiveness();
 
@@ -217,10 +217,10 @@ struct CoalesceLocals : public WalkerPass<CFGWalker<CoalesceLocals, Visitor<Coal
   }
 };
 
-void CoalesceLocals::walk(Expression*& root) {
-  numLocals = getFunction()->getNumLocals();
+void CoalesceLocals::doWalkFunction(Function* func) {
+  numLocals = func->getNumLocals();
   // collect initial liveness info
-  WalkerPass<CFGWalker<CoalesceLocals, Visitor<CoalesceLocals>, Liveness>>::walk(root);
+  WalkerPass<CFGWalker<CoalesceLocals, Visitor<CoalesceLocals>, Liveness>>::doWalkFunction(func);
   // ignore links to dead blocks, so they don't confuse us and we can see their stores are all ineffective
   liveBlocks = findLiveBlocks();
   unlinkDeadBlocks(liveBlocks);
@@ -243,7 +243,7 @@ void CoalesceLocals::walk(Expression*& root) {
   std::vector<Index> indices;
   pickIndices(indices);
   // apply indices
-  applyIndices(indices, root);
+  applyIndices(indices, func->body);
 }
 
 void CoalesceLocals::flowLiveness() {
diff --git a/src/passes/DuplicateFunctionElimination.cpp b/src/passes/DuplicateFunctionElimination.cpp
index 593ddb7d9..c33baedd4 100644
--- a/src/passes/DuplicateFunctionElimination.cpp
+++ b/src/passes/DuplicateFunctionElimination.cpp
@@ -39,16 +39,15 @@ struct FunctionHasher : public PostWalker<FunctionHasher, Visitor<FunctionHasher
     output = output_;
   }
 
-  void walk(Expression*& root) {
+  void doWalkFunction(Function* func) {
     assert(digest == 0);
-    auto* func = getFunction();
     hash(func->getNumParams());
     for (auto type : func->params) hash(type);
     hash(func->getNumVars());
     for (auto type : func->vars) hash(type);
     hash(func->result);
     hash64(func->type.is() ? uint64_t(func->type.str) : uint64_t(0));
-    hash(ExpressionAnalyzer::hash(root));
+    hash(ExpressionAnalyzer::hash(func->body));
     output->at(func) = digest;
   }
 
@@ -98,7 +97,7 @@ struct DuplicateFunctionElimination : public Pass {
       }
       FunctionHasher hasher;
       hasher.setOutput(&hashes);
-      hasher.startWalk(module);
+      hasher.walkModule(module);
       // Find hash-equal groups
       std::map<uint32_t, std::vector<Function*>> hashGroups;
       for (auto& func : module->functions) {
@@ -130,7 +129,7 @@ struct DuplicateFunctionElimination : public Pass {
         // replace direct calls
         FunctionReplacer replacer;
         replacer.setReplacements(&replacements);
-        replacer.startWalk(module);
+        replacer.walkModule(module);
         // replace in table
         for (auto& name : module->table.names) {
           auto iter = replacements.find(name);
diff --git a/src/passes/RemoveUnusedBrs.cpp b/src/passes/RemoveUnusedBrs.cpp
index 785e09570..9aa063f60 100644
--- a/src/passes/RemoveUnusedBrs.cpp
+++ b/src/passes/RemoveUnusedBrs.cpp
@@ -154,12 +154,12 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<R
     }
   }
 
-  void walk(Expression*& root) {
+  void doWalkFunction(Function* func) {
     // multiple cycles may be needed
     bool worked = false;
     do {
       anotherCycle = false;
-      WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<RemoveUnusedBrs>>>::walk(root);
+      WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<RemoveUnusedBrs>>>::doWalkFunction(func);
       assert(ifStack.empty());
       // flows may contain returns, which are flowing out and so can be optimized
       for (size_t i = 0; i < flows.size(); i++) {
@@ -203,7 +203,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<R
     };
     Selectifier selectifier;
     selectifier.setModule(getModule());
-    selectifier.walk(root);
+    selectifier.walkFunction(func);
     if (worked) {
       // Our work may alter block and if types, they may now return
       struct TypeUpdater : public WalkerPass<PostWalker<TypeUpdater, Visitor<TypeUpdater>>> {
@@ -218,7 +218,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs, Visitor<R
         }
       };
       TypeUpdater typeUpdater;
-      typeUpdater.walk(root);
+      typeUpdater.walkFunction(func);
     }
   }
 };
diff --git a/src/passes/SimplifyLocals.cpp b/src/passes/SimplifyLocals.cpp
index dac6cfa65..57057e8f9 100644
--- a/src/passes/SimplifyLocals.cpp
+++ b/src/passes/SimplifyLocals.cpp
@@ -376,7 +376,7 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals,
     self->pushTask(visitPre, currp);
   }
 
-  void walk(Expression*& root) {
+  void doWalkFunction(Function* func) {
     // multiple passes may be required per function, consider this:
     //    x = load
     //    y = store
@@ -385,7 +385,7 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals,
     do {
       anotherCycle = false;
       // main operation
-      WalkerPass<LinearExecutionWalker<SimplifyLocals, Visitor<SimplifyLocals>>>::walk(root);
+      WalkerPass<LinearExecutionWalker<SimplifyLocals, Visitor<SimplifyLocals>>>::doWalkFunction(func);
       // enlarge blocks that were marked, for the next round
       if (blocksToEnlarge.size() > 0) {
         for (auto* block : blocksToEnlarge) {
@@ -421,14 +421,14 @@ struct SimplifyLocals : public WalkerPass<LinearExecutionWalker<SimplifyLocals,
     // remove the set.
     // First, count get_locals
     std::vector<int> numGetLocals; // local => # of get_locals for it
-    numGetLocals.resize(getFunction()->getNumLocals());
+    numGetLocals.resize(func->getNumLocals());
     GetLocalCounter counter;
     counter.numGetLocals = &numGetLocals;
-    counter.walk(root);
+    counter.walkFunction(func);
     // Second, remove unneeded sets
     SetLocalRemover remover;
     remover.numGetLocals = &numGetLocals;
-    remover.walk(root);
+    remover.walkFunction(func);
   }
 };
 
diff --git a/src/wasm-linker.cpp b/src/wasm-linker.cpp
index 066ed355a..aa098c2e7 100644
--- a/src/wasm-linker.cpp
+++ b/src/wasm-linker.cpp
@@ -316,7 +316,7 @@ void Linker::emscriptenGlue(std::ostream& o) {
   };
   AsmConstWalker walker;
   walker.parent = this;
-  walker.startWalk(&out.wasm);
+  walker.walkModule(&out.wasm);
   // print
   o << "\"asmConsts\": {";
   bool first = true;
diff --git a/src/wasm-traversal.h b/src/wasm-traversal.h
index d7fc83f59..9ed02027e 100644
--- a/src/wasm-traversal.h
+++ b/src/wasm-traversal.h
@@ -184,9 +184,23 @@ struct Walker : public VisitorType {
 
   // Walk starting
 
-  void startWalk(Module *module) {
+  void walkFunction(Function* func) {
+    setFunction(func);
+    static_cast<SubType*>(this)->doWalkFunction(func);
+  }
+
+  // override this to provide custom functionality
+  void doWalkFunction(Function* func) {
+    walk(func->body);
+  }
+
+  void walkModule(Module *module) {
     setModule(module);
+    static_cast<SubType*>(this)->doWalkModule(module);
+  }
 
+  // override this to provide custom functionality
+  void doWalkModule(Module *module) {
     // Dispatch statically through the SubType.
     SubType* self = static_cast<SubType*>(this);
     for (auto& curr : module->functionTypes) {
@@ -208,8 +222,7 @@ struct Walker : public VisitorType {
         instance->setModule(module);
         allocated = std::unique_ptr<SubType>(instance);
       }
-      instance->setFunction(func);
-      instance->walk(func->body);
+      instance->walkFunction(func);
       instance->visitFunction(func);
     };
 
diff --git a/src/wasm-validator.h b/src/wasm-validator.h
index a19071959..39eaca572 100644
--- a/src/wasm-validator.h
+++ b/src/wasm-validator.h
@@ -35,7 +35,7 @@ struct WasmValidator : public PostWalker<WasmValidator, Visitor<WasmValidator>>
 public:
   bool validate(Module& module) {
     valid = true;
-    startWalk(&module);
+    walkModule(&module);
     return valid;
   }
 
@@ -280,9 +280,8 @@ public:
     }
   }
 
-  void walk(Expression*& root) {
-    //std::cerr << "start a function " << getFunction()->name << "\n";
-    PostWalker<WasmValidator, Visitor<WasmValidator>>::walk(root);
+  void doWalkFunction(Function* func) {
+    PostWalker<WasmValidator, Visitor<WasmValidator>>::doWalkFunction(func);
     shouldBeTrue(breakTypes.size() == 0, "break targets", "all break targets must be valid");
   }