Add function rename utility (#1805)

And use it in wasm-emscripten

2 files changed, 48 insertions, 21 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index d6ee50dd0..851b926c6 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -18,6 +18,7 @@
 #define wasm_ir_module_h
 
 #include "wasm.h"
+#include "ir/find_all.h"
 #include "ir/manipulation.h"
 
 namespace wasm {
@@ -128,6 +129,52 @@ inline void copyModule(Module& in, Module& out) {
   out.debugInfoFileNames = in.debugInfoFileNames;
 }
 
+// Renaming
+
+template<typename T>
+inline void renameFunctions(Module& wasm, T& map) {
+  // Update the function itself.
+  for (auto& pair : map) {
+    assert(!wasm.getFunctionOrNull(pair.second));
+    wasm.getFunction(pair.first)->name = pair.second;
+  }
+  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);
+  for (auto& segment : wasm.table.segments) {
+    for (auto& name : segment.data) {
+      maybeUpdate(name);
+    }
+  }
+  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);
+      }
+    }
+  }
+}
+
+inline void renameFunction(Module& wasm, Name oldName, Name newName) {
+  std::map<Name, Name> map;
+  map[oldName] = newName;
+  renameFunctions(wasm, map);
+}
+
 // Convenient iteration over imported/non-imported module elements
 
 template<typename T>
diff --git a/src/wasm/wasm-emscripten.cpp b/src/wasm/wasm-emscripten.cpp
index 92a11b8f1..d96a6636e 100644
--- a/src/wasm/wasm-emscripten.cpp
+++ b/src/wasm/wasm-emscripten.cpp
@@ -706,31 +706,11 @@ struct FixInvokeFunctionNamesWalker : public PostWalker<FixInvokeFunctionNamesWa
     }
   }
 
-  void visitTable(Table* curr) {
-    for (auto& segment : curr->segments) {
-      for (size_t i = 0; i < segment.data.size(); i++) {
-        auto it = importRenames.find(segment.data[i]);
-        if (it != importRenames.end()) {
-          segment.data[i] = it->second;
-        }
-      }
-    }
-  }
-
-  void visitCall(Call* curr) {
-    if (wasm.getFunction(curr->target)->imported()) {
-      auto it = importRenames.find(curr->target);
-      if (it != importRenames.end()) {
-        curr->target = it->second;
-      }
-    }
-  }
-
   void visitModule(Module* curr) {
     for (auto importName : toRemove) {
       wasm.removeFunction(importName);
     }
-    wasm.updateMaps();
+    ModuleUtils::renameFunctions(wasm, importRenames);
   }
 };