[NFC] Simplify traversal code for setting the module (#5082)

Make walkModuleCode set the module automatically, like walkModule already does.
Also remove some unneeded module settings when calling those methods.

2 files changed, 3 insertions, 4 deletions

diff --git a/src/pass.h b/src/pass.h
index 623c1f971..45780f9e5 100644
--- a/src/pass.h
+++ b/src/pass.h
@@ -460,20 +460,17 @@ public:
     }
     // Single-thread running just calls the walkModule traversal.
     setPassRunner(runner);
-    WalkerType::setModule(module);
     WalkerType::walkModule(module);
   }
 
   void
   runOnFunction(PassRunner* runner, Module* module, Function* func) override {
     setPassRunner(runner);
-    WalkerType::setModule(module);
-    WalkerType::walkFunction(func);
+    WalkerType::walkFunctionInModule(func, module);
   }
 
   void runOnModuleCode(PassRunner* runner, Module* module) {
     setPassRunner(runner);
-    WalkerType::setModule(module);
     WalkerType::walkModuleCode(module);
   }
 
diff --git a/src/wasm-traversal.h b/src/wasm-traversal.h
index a77356a4c..34e1f9453 100644
--- a/src/wasm-traversal.h
+++ b/src/wasm-traversal.h
@@ -268,6 +268,7 @@ struct Walker : public VisitorType {
 
   // Walks module-level code, that is, code that is not in functions.
   void walkModuleCode(Module* module) {
+    setModule(module);
     // Dispatch statically through the SubType.
     SubType* self = static_cast<SubType*>(this);
     for (auto& curr : module->globals) {
@@ -283,6 +284,7 @@ struct Walker : public VisitorType {
         self->walk(item);
       }
     }
+    setModule(nullptr);
   }
 
   // Walk implementation. We don't use recursion as ASTs may be highly