5 files changed, 147 insertions, 89 deletions
diff --git a/src/s2wasm.h b/src/s2wasm.h
index 0efe0408e..3cb4c7139 100644
--- a/src/s2wasm.h
+++ b/src/s2wasm.h
@@ -253,6 +253,7 @@ class S2WasmBuilder {
     return new LinkerObject::Relocation(
         kind, target, fixEmLongjmp(cleanFunction(name)), offset);
   }
+
   Expression* relocationToGetGlobal(LinkerObject::Relocation* relocation) {
     if (!relocation) {
       return nullptr;
@@ -421,6 +422,7 @@ class S2WasmBuilder {
   bool isFunctionName(Name name) {
     return !!strstr(name.str, "@FUNCTION");
   }
+
   // Drop the @ and after it.
   Name cleanFunction(Name name) {
     if (!strchr(name.str, '@')) return name;
@@ -615,7 +617,7 @@ class S2WasmBuilder {
 
   void parseFuncType() {
     auto decl = make_unique<FunctionType>();
-    Name rawName = getCommaSeparated();
+    Name name = getCommaSeparated();
     skipComma();
     if(match("void")) {
       decl->result = none;
@@ -627,7 +629,6 @@ class S2WasmBuilder {
     decl->name = "FUNCSIG$" + sig;
 
     FunctionType *ty = wasm->getFunctionTypeOrNull(decl->name);
-    Name name = fixEmEHSjLjNames(rawName, sig);
     if (!ty) {
       // The wasm module takes ownership of the FunctionType if we insert it.
       // Otherwise it's already in the module and ours is freed.
@@ -979,14 +980,11 @@ class S2WasmBuilder {
         if (*s == ',') {
           skipComma();
           int num = getNumInputs();
-          auto inputs = getInputs(num);
-          for (int i = 0; i < num; i++) {
-            curr->operands.push_back(inputs[i]);
+          for (Expression* input : getInputs(num)) {
+            curr->operands.push_back(input);
           }
         }
-        Name target = linkerObj->resolveAlias(
-            fixEmEHSjLjNames(rawTarget, curr->type, curr->operands),
-            LinkerObject::Relocation::kFunction);
+        Name target = linkerObj->resolveAlias(rawTarget, LinkerObject::Relocation::kFunction);
         curr->target = target;
         if (!linkerObj->isFunctionImplemented(target)) {
           linkerObj->addUndefinedFunctionCall(curr);
@@ -1447,83 +1445,14 @@ class S2WasmBuilder {
     }
   }
 
-  // Fixes function name hacks caused by LLVM exception & setjmp/longjmp
-  // handling pass for wasm.
-  // This does two things:
-  // 1. Change emscripten_longjmp_jmpbuf to emscripten_longjmp.
-  //    In setjmp/longjmp handling pass in wasm backend, what we want to do is
-  //    to change all function calls to longjmp to calls to emscripten_longjmp.
-  //    Because we replace all calls to longjmp to emscripten_longjmp, the
-  //    signature of that function should be the same as longjmp:
-  //    emscripten_longjmp(jmp_buf, int)
-  //    But after calling a function that might longjmp, while we test whether
-  //    a longjmp occurred, we have to load an int address value and call
-  //    emscripten_longjmp again with that address as the first argument. (Refer
-  //    to lib/Target/WebAssembly/WebAssemblyEmscriptenEHSjLj.cpp in LLVM for
-  //    details.)
-  //    In this case we need the signature of emscripten_longjmp to be (int,
-  //    int). So we need two different kinds of emscripten_longjmp signatures in
-  //    LLVM IR. Both signatures will be lowered to (int, int) eventually, but
-  //    in LLVM IR, types are not lowered yet.
-  //    So we declare two functions in LLVM:
-  //    emscripten_longjmp_jmpbuf(jmp_buf, int)
-  //    emscripten_longjmp(int, int)
-  //    And we change the name of emscripten_longjmp_jmpbuf to
-  //    emscripten_longjmp here.
-  // 2. Converts invoke wrapper names.
-  //    Refer to the comments in fixEmExceptionInvoke below.
-  template<typename ListType>
-  Name fixEmEHSjLjNames(const Name &name, Type result,
-                        const ListType &operands) {
-    return fixEmEHSjLjNames(name, getSig(result, operands));
-  }
-
-  Name fixEmEHSjLjNames(const Name &name, const std::string& sig) {
-    if (name == "emscripten_longjmp_jmpbuf")
-      return "emscripten_longjmp";
-    return fixEmExceptionInvoke(name, sig);
-  }
-
   // This version only converts emscripten_longjmp_jmpbuf and does not deal
   // with invoke wrappers. This is used when we only have a function name as
   // relocatable constant.
-  Name fixEmLongjmp(const Name &name) {
+  static Name fixEmLongjmp(const Name &name) {
     if (name == "emscripten_longjmp_jmpbuf")
       return "emscripten_longjmp";
     return name;
   }
-
-  // Converts invoke wrapper names generated by LLVM backend to real invoke
-  // wrapper names that are expected by JavaScript glue code.
-  // This is required to support wasm exception handling (asm.js style).
-  //
-  // LLVM backend lowers
-  //   invoke @func(arg1, arg2) to label %invoke.cont unwind label %lpad
-  // into
-  // ... (some code)
-  //   call @invoke_SIG(func, arg1, arg2)
-  // ... (some code)
-  // SIG is a mangled string generated based on the LLVM IR-level function
-  // signature. In LLVM IR, types are not lowered yet, so this mangling scheme
-  // simply takes LLVM's string representtion of parameter types and concatenate
-  // them with '_'. For example, the name of an invoke wrapper for function
-  // void foo(struct mystruct*, int) will be
-  // "__invoke_void_%struct.mystruct*_int".
-  // This function converts the names of invoke wrappers based on their lowered
-  // argument types and a return type. In the example above, the resulting new
-  // wrapper name becomes "invoke_vii".
-  Name fixEmExceptionInvoke(const Name &name, const std::string& sig) {
-    std::string nameStr = name.c_str();
-    if (nameStr.front() == '"' && nameStr.back() == '"') {
-      nameStr = nameStr.substr(1, nameStr.size() - 2);
-    }
-    if (nameStr.find("__invoke_") != 0) {
-      return name;
-    }
-    std::string sigWoOrigFunc = sig.front() + sig.substr(2, sig.size() - 2);
-    return Name("invoke_" + sigWoOrigFunc);
-  }
-
 };
 
 } // namespace wasm
diff --git a/src/tools/wasm-emscripten-finalize.cpp b/src/tools/wasm-emscripten-finalize.cpp
index 1f599efe9..7d00e8116 100644
--- a/src/tools/wasm-emscripten-finalize.cpp
+++ b/src/tools/wasm-emscripten-finalize.cpp
@@ -108,6 +108,7 @@ int main(int argc, const char *argv[]) {
   initializerFunctions.push_back("__wasm_call_ctors");
 
   EmscriptenGlueGenerator generator(wasm);
+  generator.fixInvokeFunctionNames();
   generator.generateRuntimeFunctions();
   generator.generateMemoryGrowthFunction();
   generator.generateDynCallThunks();
diff --git a/src/wasm-emscripten.h b/src/wasm-emscripten.h
index d2a0a8e3e..a0a87664b 100644
--- a/src/wasm-emscripten.h
+++ b/src/wasm-emscripten.h
@@ -51,6 +51,8 @@ public:
   // Replace placeholder emscripten_asm_const functions with *_signature versions.
   void fixEmAsmConsts();
 
+  void fixInvokeFunctionNames();
+
 private:
   Module& wasm;
   Builder builder;
diff --git a/src/wasm/wasm-emscripten.cpp b/src/wasm/wasm-emscripten.cpp
index 9a393db43..5b5d557e7 100644
--- a/src/wasm/wasm-emscripten.cpp
+++ b/src/wasm/wasm-emscripten.cpp
@@ -508,6 +508,128 @@ void EmscriptenGlueGenerator::fixEmAsmConsts() {
   fixEmJsFuncsAndReturnWalker(wasm);
 }
 
+
+// Fixes function name hacks caused by LLVM exception & setjmp/longjmp
+// handling pass for wasm.
+// This does two things:
+// 1. Change emscripten_longjmp_jmpbuf to emscripten_longjmp.
+//    In setjmp/longjmp handling pass in wasm backend, what we want to do is
+//    to change all function calls to longjmp to calls to emscripten_longjmp.
+//    Because we replace all calls to longjmp to emscripten_longjmp, the
+//    signature of that function should be the same as longjmp:
+//    emscripten_longjmp(jmp_buf, int)
+//    But after calling a function that might longjmp, while we test whether
+//    a longjmp occurred, we have to load an int address value and call
+//    emscripten_longjmp again with that address as the first argument. (Refer
+//    to lib/Target/WebAssembly/WebAssemblyEmscriptenEHSjLj.cpp in LLVM for
+//    details.)
+//    In this case we need the signature of emscripten_longjmp to be (int,
+//    int). So we need two different kinds of emscripten_longjmp signatures in
+//    LLVM IR. Both signatures will be lowered to (int, int) eventually, but
+//    in LLVM IR, types are not lowered yet.
+//    So we declare two functions in LLVM:
+//    emscripten_longjmp_jmpbuf(jmp_buf, int)
+//    emscripten_longjmp(int, int)
+//    And we change the name of emscripten_longjmp_jmpbuf to
+//    emscripten_longjmp here.
+// 2. Converts invoke wrapper names.
+//    Refer to the comments in fixEmExceptionInvoke below.
+struct FixInvokeFunctionNamesWalker : public PostWalker<FixInvokeFunctionNamesWalker> {
+  Module& wasm;
+  std::map<Name, Name> importRenames;
+  std::vector<Name> toRemove;
+  std::set<Name> newImports;
+
+  FixInvokeFunctionNamesWalker(Module& _wasm) : wasm(_wasm) {}
+
+  // Converts invoke wrapper names generated by LLVM backend to real invoke
+  // wrapper names that are expected by JavaScript glue code.
+  // This is required to support wasm exception handling (asm.js style).
+  //
+  // LLVM backend lowers
+  //   invoke @func(arg1, arg2) to label %invoke.cont unwind label %lpad
+  // into
+  // ... (some code)
+  //   call @invoke_SIG(func, arg1, arg2)
+  // ... (some code)
+  // SIG is a mangled string generated based on the LLVM IR-level function
+  // signature. In LLVM IR, types are not lowered yet, so this mangling scheme
+  // simply takes LLVM's string representtion of parameter types and concatenate
+  // them with '_'. For example, the name of an invoke wrapper for function
+  // void foo(struct mystruct*, int) will be
+  // "__invoke_void_%struct.mystruct*_int".
+  // This function converts the names of invoke wrappers based on their lowered
+  // argument types and a return type. In the example above, the resulting new
+  // wrapper name becomes "invoke_vii".
+  static Name fixEmExceptionInvoke(const Name& name, const std::string& sig) {
+    std::string nameStr = name.c_str();
+    if (nameStr.front() == '"' && nameStr.back() == '"') {
+      nameStr = nameStr.substr(1, nameStr.size() - 2);
+    }
+    if (nameStr.find("__invoke_") != 0) {
+      return name;
+    }
+    std::string sigWoOrigFunc = sig.front() + sig.substr(2, sig.size() - 2);
+    return Name("invoke_" + sigWoOrigFunc);
+  }
+
+  static Name fixEmEHSjLjNames(const Name &name, const std::string& sig) {
+    if (name == "emscripten_longjmp_jmpbuf")
+      return "emscripten_longjmp";
+    return fixEmExceptionInvoke(name, sig);
+  }
+
+  void visitImport(Import* curr) {
+    if (curr->kind != ExternalKind::Function)
+      return;
+
+    FunctionType* func = wasm.getFunctionType(curr->functionType);
+    Name newname = fixEmEHSjLjNames(curr->name, getSig(func));
+    if (newname == curr->name)
+      return;
+
+    assert(importRenames.count(curr->name) == 0);
+    importRenames[curr->name] = newname;
+    // Either rename of remove the existing import
+    if (wasm.getImportOrNull(newname) || !newImports.insert(newname).second) {
+      toRemove.push_back(curr->name);
+    } else {
+      curr->base = newname;
+      curr->name = newname;
+    }
+  }
+
+  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 visitCallImport(CallImport* curr) {
+    auto it = importRenames.find(curr->target);
+    if (it != importRenames.end()) {
+      curr->target = it->second;
+    }
+  }
+
+  void visitModule(Module* curr) {
+    for (auto importName : toRemove) {
+      wasm.removeImport(importName);
+    }
+    wasm.updateMaps();
+  }
+};
+
+void EmscriptenGlueGenerator::fixInvokeFunctionNames() {
+  FixInvokeFunctionNamesWalker walker(wasm);
+  walker.walkModule(&wasm);
+}
+
 template<class C>
 void printSet(std::ostream& o, C& c) {
   o << "[";
@@ -587,15 +709,18 @@ std::string EmscriptenGlueGenerator::generateEmscriptenMetadata(
     meta << "]";
   }
 
+  // We use the `base` rather than the `name` of the imports here and below
+  // becasue this is the externally visible name that the embedder (JS) will
+  // see.
   meta << ", \"declares\": [";
   commaFirst = true;
   for (const auto& import : wasm.imports) {
     if (import->kind == ExternalKind::Function &&
-        (emJsWalker.codeByName.count(import->name.str) == 0) &&
-        !import->name.startsWith(EMSCRIPTEN_ASM_CONST.str) &&
-        !import->name.startsWith("invoke_") &&
-        !import->name.startsWith("jsCall_")) {
-      meta << maybeComma() << '"' << import->name.str << '"';
+        (emJsWalker.codeByName.count(import->base.str) == 0) &&
+        !import->base.startsWith(EMSCRIPTEN_ASM_CONST.str) &&
+        !import->base.startsWith("invoke_") &&
+        !import->base.startsWith("jsCall_")) {
+      meta << maybeComma() << '"' << import->base.str << '"';
     }
   }
   meta << "]";
@@ -604,7 +729,7 @@ std::string EmscriptenGlueGenerator::generateEmscriptenMetadata(
   commaFirst = true;
   for (const auto& import : wasm.imports) {
     if (import->kind == ExternalKind::Global) {
-      meta << maybeComma() << "\"_" << import->name.str << '"';
+      meta << maybeComma() << "\"_" << import->base.str << '"';
     }
   }
   meta << "]";
@@ -626,8 +751,8 @@ std::string EmscriptenGlueGenerator::generateEmscriptenMetadata(
   meta << ", \"invokeFuncs\": [";
   commaFirst = true;
   for (const auto& import : wasm.imports) {
-    if (import->name.startsWith("invoke_")) {
-      meta << maybeComma() << '"' << import->name.str << '"';
+    if (import->base.startsWith("invoke_")) {
+      meta << maybeComma() << '"' << import->base.str << '"';
     }
   }
   meta << "]";
@@ -645,6 +770,7 @@ std::string emscriptenGlue(
     std::vector<Name> const& initializerFunctions,
     unsigned numReservedFunctionPointers) {
   EmscriptenGlueGenerator generator(wasm, stackPointer);
+  generator.fixInvokeFunctionNames();
   generator.generateRuntimeFunctions();
 
   if (allowMemoryGrowth) {
diff --git a/test/dot_s/fix_em_ehsjlj_names.wast b/test/dot_s/fix_em_ehsjlj_names.wast
index 3a46b8b74..3fbac3936 100644
--- a/test/dot_s/fix_em_ehsjlj_names.wast
+++ b/test/dot_s/fix_em_ehsjlj_names.wast
@@ -7,11 +7,11 @@
  (type $FUNCSIG$v (func))
  (type $FUNCSIG$ffd (func (param f32 f64) (result f32)))
  (type $FUNCSIG$iii (func (param i32 i32) (result i32)))
- (import "env" "emscripten_longjmp" (func $emscripten_longjmp (param i32 i32)))
- (import "env" "invoke_ffd" (func $invoke_ffd (param i32 f32 f64) (result f32)))
  (import "env" "invoke_iii" (func $invoke_iii (param i32 i32 i32) (result i32)))
+ (import "env" "invoke_ffd" (func $invoke_ffd (param i32 f32 f64) (result f32)))
  (import "env" "invoke_iiii" (func $invoke_iiii (param i32 i32 i32 i32) (result i32)))
  (import "env" "invoke_v" (func $invoke_v (param i32)))
+ (import "env" "emscripten_longjmp" (func $emscripten_longjmp (param i32 i32)))
  (import "env" "memory" (memory $0 1))
  (table 5 5 anyfunc)
  (elem (i32.const 0) $__wasm_nullptr $_Z5func1v $_Z5func2iii $_Z5func3fd $_Z5func4P8mystructS_)
@@ -136,4 +136,4 @@
   )
  )
 )
-;; METADATA: { "asmConsts": {},"staticBump": 12, "initializers": [], "declares": ["emscripten_longjmp"], "externs": [], "implementedFunctions": ["__Z5func1v","__Z5func2iii","__Z5func3fd","__Z5func4P8mystructS_","_main","___wasm_nullptr","_stackSave","_stackAlloc","_stackRestore","_dynCall_v","_dynCall_iiii","_dynCall_ffd","_dynCall_iii"], "exports": ["main","stackSave","stackAlloc","stackRestore","dynCall_v","dynCall_iiii","dynCall_ffd","dynCall_iii"], "invokeFuncs": ["invoke_ffd","invoke_iii","invoke_iiii","invoke_v"] }
+;; METADATA: { "asmConsts": {},"staticBump": 12, "initializers": [], "declares": ["emscripten_longjmp"], "externs": [], "implementedFunctions": ["__Z5func1v","__Z5func2iii","__Z5func3fd","__Z5func4P8mystructS_","_main","___wasm_nullptr","_stackSave","_stackAlloc","_stackRestore","_dynCall_v","_dynCall_iiii","_dynCall_ffd","_dynCall_iii"], "exports": ["main","stackSave","stackAlloc","stackRestore","dynCall_v","dynCall_iiii","dynCall_ffd","dynCall_iii"], "invokeFuncs": ["invoke_iii","invoke_ffd","invoke_iiii","invoke_v"] }