Emscripten addFunction support for Wasm backend (#1395)

This adds necessary command line options for addFunction support, and generates required jsCall imports and generates jsCall thunk functions.

5 files changed, 168 insertions, 11 deletions

diff --git a/src/tools/s2wasm.cpp b/src/tools/s2wasm.cpp
index d4e7731c5..a2c4e4a38 100644
--- a/src/tools/s2wasm.cpp
+++ b/src/tools/s2wasm.cpp
@@ -41,6 +41,7 @@ int main(int argc, const char *argv[]) {
   std::string startFunction;
   std::vector<std::string> archiveLibraries;
   TrapMode trapMode = TrapMode::Allow;
+  unsigned numReservedFunctionPointers = 0;
   Options options("s2wasm", "Link .s file into .wast");
   options.extra["validate"] = "wasm";
   options
@@ -121,6 +122,14 @@ int main(int argc, const char *argv[]) {
              }
              o->extra["validate"] = argument;
            })
+      .add("--emscripten-reserved-function-pointers", "",
+           "Number of reserved function pointers for emscripten addFunction "
+           "support",
+           Options::Arguments::One,
+           [&numReservedFunctionPointers](Options *o,
+                                          const std::string &argument) {
+             numReservedFunctionPointers = std::stoi(argument);
+           })
       .add_positional("INFILE", Options::Arguments::One,
                       [](Options *o, const std::string& argument) {
                         o->extra["infile"] = argument;
@@ -189,7 +198,8 @@ int main(int argc, const char *argv[]) {
       allowMemoryGrowth,
       linker.getStackPointerAddress(),
       linker.getStaticBump(),
-      linker.getOutput().getInitializerFunctions());
+      linker.getOutput().getInitializerFunctions(),
+      numReservedFunctionPointers);
   }
 
   if (options.extra["validate"] != "none") {
diff --git a/src/tools/wasm-emscripten-finalize.cpp b/src/tools/wasm-emscripten-finalize.cpp
index 223fa3082..28f5a7c06 100644
--- a/src/tools/wasm-emscripten-finalize.cpp
+++ b/src/tools/wasm-emscripten-finalize.cpp
@@ -39,6 +39,7 @@ int main(int argc, const char *argv[]) {
   std::string infile;
   std::string outfile;
   bool emitBinary = true;
+  unsigned numReservedFunctionPointers = 0;
   std::vector<Name> forcedExports;
   Options options("wasm-emscripten-finalize",
                   "Performs Emscripten-specific transforms on .wasm files");
@@ -54,6 +55,14 @@ int main(int argc, const char *argv[]) {
            [&emitBinary](Options*, const std::string& ) {
              emitBinary = false;
            })
+      .add("--emscripten-reserved-function-pointers", "",
+           "Number of reserved function pointers for emscripten addFunction "
+           "support",
+           Options::Arguments::One,
+           [&numReservedFunctionPointers](Options *,
+                                          const std::string &argument) {
+             numReservedFunctionPointers = std::stoi(argument);
+           })
       .add_positional("INFILE", Options::Arguments::One,
                       [&infile](Options *o, const std::string& argument) {
                         infile = argument;
@@ -80,6 +89,7 @@ int main(int argc, const char *argv[]) {
   generator.generateRuntimeFunctions();
   generator.generateMemoryGrowthFunction();
   generator.generateDynCallThunks();
+  generator.generateJSCallThunks(numReservedFunctionPointers);
   generator.fixEmAsmConsts();
 
   if (options.debug) {
diff --git a/src/tools/wasm-link-metadata.cpp b/src/tools/wasm-link-metadata.cpp
index 96e06f95b..3af00fbbf 100644
--- a/src/tools/wasm-link-metadata.cpp
+++ b/src/tools/wasm-link-metadata.cpp
@@ -67,6 +67,7 @@ void parseLinkingSection(std::vector<char> const& data, uint32_t &dataSize) {
 int main(int argc, const char *argv[]) {
   std::string infile;
   std::string outfile;
+  unsigned numReservedFunctionPointers = 0;
   Options options("wasm-link-metadata",
                   "Reads wasm .o file and emits .json metadata");
   options
@@ -76,6 +77,14 @@ int main(int argc, const char *argv[]) {
              outfile = argument;
              Colors::disable();
            })
+      .add("--emscripten-reserved-function-pointers", "",
+           "Number of reserved function pointers for emscripten addFunction "
+           "support",
+           Options::Arguments::One,
+           [&numReservedFunctionPointers](Options *o,
+                                          const std::string &argument) {
+             numReservedFunctionPointers = std::stoi(argument);
+           })
       .add_positional("INFILE", Options::Arguments::One,
                       [&infile](Options *o, const std::string& argument) {
                         infile = argument;
@@ -110,7 +119,8 @@ int main(int argc, const char *argv[]) {
   initializerFunctions.push_back("__wasm_call_ctors");
 
   EmscriptenGlueGenerator generator(wasm);
-  std::string metadata = generator.generateEmscriptenMetadata(dataSize, initializerFunctions);
+  std::string metadata = generator.generateEmscriptenMetadata(
+      dataSize, initializerFunctions, numReservedFunctionPointers);
   Output output(outfile, Flags::Text, Flags::Release);
   output << metadata;
 
diff --git a/src/wasm-emscripten.cpp b/src/wasm-emscripten.cpp
index 60bf5cce3..821f3aa9d 100644
--- a/src/wasm-emscripten.cpp
+++ b/src/wasm-emscripten.cpp
@@ -196,6 +196,106 @@ void EmscriptenGlueGenerator::generateDynCallThunks() {
   }
 }
 
+struct JSCallWalker : public PostWalker<JSCallWalker> {
+  Module &wasm;
+  JSCallWalker(Module &_wasm) : wasm(_wasm) {
+    if (wasm.table.segments.size() == 0) {
+      auto emptySegment =
+          wasm.allocator.alloc<Const>()->set(Literal(uint32_t(0)));
+      wasm.table.segments.emplace_back(emptySegment);
+    }
+    const auto& tableSegmentData = wasm.table.segments[0].data;
+
+    // Check if jsCalls have already been created
+    for (Index i = 0; i < tableSegmentData.size(); ++i) {
+      if (tableSegmentData[i].startsWith("jsCall_")) {
+        jsCallStartIndex = i;
+        return;
+      }
+    }
+    jsCallStartIndex =
+        wasm.table.segments[0].offset->cast<Const>()->value.getInteger() +
+        tableSegmentData.size();
+  }
+
+  // Gather all function signatures used in call_indirect, because any of them
+  // can be used to call function pointers created by emscripten's addFunction.
+  void visitCallIndirect(CallIndirect *curr) {
+    // dynCall thunks are generated in binaryen and call_indirect instructions
+    // within them cannot be used to call function pointers returned by
+    // emscripten's addFunction.
+    if (!getFunction()->name.startsWith("dynCall_")) {
+      indirectlyCallableSigs.insert(
+          getSig(wasm.getFunctionType(curr->fullType)));
+    }
+  }
+
+  bool createJSCallThunks;
+  Index jsCallStartIndex;
+  // Function type signatures used in call_indirect instructions
+  std::set<std::string> indirectlyCallableSigs;
+};
+
+JSCallWalker getJSCallWalker(Module& wasm) {
+  JSCallWalker walker(wasm);
+  walker.walkModule(&wasm);
+  return walker;
+}
+
+void EmscriptenGlueGenerator::generateJSCallThunks(
+    unsigned numReservedFunctionPointers) {
+  if (numReservedFunctionPointers == 0)
+    return;
+
+  JSCallWalker walker = getJSCallWalker(wasm);
+  auto& tableSegmentData = wasm.table.segments[0].data;
+  for (std::string sig : walker.indirectlyCallableSigs) {
+    // Add imports for jsCall_sig (e.g. jsCall_vi).
+    // Imported jsCall_sig functions have their first parameter as an index to
+    // the function table, so we should prepend an 'i' to parameters' signature
+    // (e.g. If the signature of the callee is 'vi', the imported jsCall_vi
+    // function would have signature 'vii'.)
+    std::string importSig = std::string(1, sig[0]) + 'i' + sig.substr(1);
+    FunctionType *importType = ensureFunctionType(importSig, &wasm);
+    auto import = new Import;
+    import->name = import->base = "jsCall_" + sig;
+    import->module = ENV;
+    import->functionType = importType->name;
+    import->kind = ExternalKind::Function;
+    wasm.addImport(import);
+    FunctionType *funcType = ensureFunctionType(sig, &wasm);
+
+    // Create jsCall_sig_index thunks (e.g. jsCall_vi_0, jsCall_vi_1, ...)
+    // e.g. If # of reserved function pointers (given by a command line
+    // argument) is 3 and there are two possible signature 'vi' and 'ii', the
+    // genereated thunks will be jsCall_vi_0, jsCall_vi_1, jsCall_vi_2,
+    // jsCall_ii_0, jsCall_ii_1, and jsCall_ii_2.
+    for (unsigned fp = 0; fp < numReservedFunctionPointers; ++fp) {
+      std::vector<NameType> params;
+      int p = 0;
+      for (const auto& ty : funcType->params) {
+        params.emplace_back(std::to_string(p++), ty);
+      }
+      Function* f = builder.makeFunction(
+          std::string("jsCall_") + sig + "_" + std::to_string(fp),
+          std::move(params), funcType->result, {});
+      std::vector<Expression*> args;
+      args.push_back(builder.makeConst(Literal(fp)));
+      for (unsigned i = 0; i < funcType->params.size(); ++i) {
+        args.push_back(builder.makeGetLocal(i, funcType->params[i]));
+      }
+      Expression* call =
+          builder.makeCallImport(import->name, args, funcType->result);
+      f->body = call;
+      wasm.addFunction(f);
+      tableSegmentData.push_back(f->name);
+    }
+  }
+  wasm.table.initial = wasm.table.max =
+      wasm.table.segments[0].offset->cast<Const>()->value.getInteger() +
+      tableSegmentData.size();
+}
+
 struct AsmConstWalker : public PostWalker<AsmConstWalker> {
   Module& wasm;
   std::vector<Address> segmentOffsets; // segment index => address offset
@@ -362,22 +462,22 @@ void printSet(std::ostream& o, C& c) {
 }
 
 std::string EmscriptenGlueGenerator::generateEmscriptenMetadata(
-    Address staticBump,
-    std::vector<Name> const& initializerFunctions) {
+    Address staticBump, std::vector<Name> const& initializerFunctions,
+    unsigned numReservedFunctionPointers) {
   std::stringstream meta;
   meta << "{ ";
 
-  AsmConstWalker walker = fixEmAsmConstsAndReturnWalker(wasm);
+  AsmConstWalker emAsmWalker = fixEmAsmConstsAndReturnWalker(wasm);
 
   // print
   meta << "\"asmConsts\": {";
   bool first = true;
-  for (auto& pair : walker.sigsForCode) {
+  for (auto& pair : emAsmWalker.sigsForCode) {
     auto& code = pair.first;
     auto& sigs = pair.second;
     if (first) first = false;
     else meta << ",";
-    meta << '"' << walker.ids[code] << "\": [\"" << code << "\", ";
+    meta << '"' << emAsmWalker.ids[code] << "\": [\"" << code << "\", ";
     printSet(meta, sigs);
     meta << "]";
   }
@@ -394,6 +494,20 @@ std::string EmscriptenGlueGenerator::generateEmscriptenMetadata(
   }
   meta << "]";
 
+  if (numReservedFunctionPointers) {
+    JSCallWalker jsCallWalker = getJSCallWalker(wasm);
+    meta << ", ";
+    meta << "\"jsCallStartIndex\": " << jsCallWalker.jsCallStartIndex << ", ";
+    meta << "\"jsCallFuncType\": [";
+    bool first = true;
+    for (std::string sig : jsCallWalker.indirectlyCallableSigs) {
+      if (!first) meta << ", ";
+      first = false;
+      meta << "\"" << sig << "\"";
+    }
+    meta << "]";
+  }
+
   meta << " }\n";
 
   return meta.str();
@@ -404,7 +518,8 @@ std::string emscriptenGlue(
     bool allowMemoryGrowth,
     Address stackPointer,
     Address staticBump,
-    std::vector<Name> const& initializerFunctions) {
+    std::vector<Name> const& initializerFunctions,
+    unsigned numReservedFunctionPointers) {
   EmscriptenGlueGenerator generator(wasm, stackPointer);
   generator.generateRuntimeFunctions();
 
@@ -414,7 +529,12 @@ std::string emscriptenGlue(
 
   generator.generateDynCallThunks();
 
-  return generator.generateEmscriptenMetadata(staticBump, initializerFunctions);
+  if (numReservedFunctionPointers) {
+    generator.generateJSCallThunks(numReservedFunctionPointers);
+  }
+
+  return generator.generateEmscriptenMetadata(staticBump, initializerFunctions,
+                                              numReservedFunctionPointers);
 }
 
 } // namespace wasm
diff --git a/src/wasm-emscripten.h b/src/wasm-emscripten.h
index 1878e6531..d2a0a8e3e 100644
--- a/src/wasm-emscripten.h
+++ b/src/wasm-emscripten.h
@@ -39,8 +39,14 @@ public:
   // signature in the indirect function table.
   void generateDynCallThunks();
 
+  // Create thunks to support emscripten's addFunction functionality. Creates (#
+  // of reserved function pointers) thunks for each indirectly called function
+  // signature.
+  void generateJSCallThunks(unsigned numReservedFunctionPointers);
+
   std::string generateEmscriptenMetadata(
-    Address staticBump, std::vector<Name> const& initializerFunctions);
+      Address staticBump, std::vector<Name> const& initializerFunctions,
+      unsigned numReservedFunctionPointers);
 
   // Replace placeholder emscripten_asm_const functions with *_signature versions.
   void fixEmAsmConsts();
@@ -64,7 +70,8 @@ std::string emscriptenGlue(
     bool allowMemoryGrowth,
     Address stackPointer,
     Address staticBump,
-    std::vector<Name> const& initializerFunctions);
+    std::vector<Name> const& initializerFunctions,
+    unsigned numReservedFunctionPointers);
 
 } // namespace wasm