1 files changed, 29 insertions, 26 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index 81f832b40..319f1a7a4 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -176,40 +176,43 @@ inline void clearModule(Module& wasm) {
 // Rename functions along with all their uses.
 // Note that for this to work the functions themselves don't necessarily need
 // to exist.  For example, it is possible to remove a given function and then
-// call this redirect all of its uses.
+// call this to redirect all of its uses.
 template<typename T> inline void renameFunctions(Module& wasm, T& map) {
   // Update the function itself.
   for (auto& [oldName, newName] : map) {
-    if (Function* F = wasm.getFunctionOrNull(oldName)) {
-      assert(!wasm.getFunctionOrNull(newName) || F->name == newName);
-      F->name = newName;
+    if (Function* func = wasm.getFunctionOrNull(oldName)) {
+      assert(!wasm.getFunctionOrNull(newName) || func->name == newName);
+      func->name = newName;
     }
   }
   wasm.updateMaps();
-  // Update other global things.
-  auto maybeUpdate = [&](Name& name) {
-    auto iter = map.find(name);
-    if (iter != map.end()) {
-      name = iter->second;
-    }
-  };
-  maybeUpdate(wasm.start);
-  ElementUtils::iterAllElementFunctionNames(&wasm, maybeUpdate);
-  for (auto& exp : wasm.exports) {
-    if (exp->kind == ExternalKind::Function) {
-      maybeUpdate(exp->value);
-    }
-  }
-  // Update call instructions.
-  for (auto& func : wasm.functions) {
-    // TODO: parallelize
-    if (!func->imported()) {
-      FindAll<Call> calls(func->body);
-      for (auto* call : calls.list) {
-        maybeUpdate(call->target);
+
+  // Update all references to it.
+  struct Updater : public WalkerPass<PostWalker<Updater>> {
+    bool isFunctionParallel() override { return true; }
+
+    T& map;
+
+    void maybeUpdate(Name& name) {
+      if (auto iter = map.find(name); iter != map.end()) {
+        name = iter->second;
       }
     }
-  }
+
+    Updater(T& map) : map(map) {}
+
+    Updater* create() override { return new Updater(map); }
+
+    void visitCall(Call* curr) { maybeUpdate(curr->target); }
+
+    void visitRefFunc(RefFunc* curr) { maybeUpdate(curr->func); }
+  };
+
+  Updater updater(map);
+  updater.maybeUpdate(wasm.start);
+  PassRunner runner(&wasm);
+  updater.run(&runner, &wasm);
+  updater.runOnModuleCode(&runner, &wasm);
 }
 
 inline void renameFunction(Module& wasm, Name oldName, Name newName) {