EM_JS binaryen support (#1410)

* Emit EM_JS metadata

* Include s2wasm-style em_js support

* Change em_js metadata to be keyed on name

* Add testcase for em_js, don't always emit emJsFuncs metadata

* Better error handling for unexpectedly-formatted __em_js__ functions

2 files changed, 137 insertions, 62 deletions

diff --git a/src/emscripten-optimizer/istring.h b/src/emscripten-optimizer/istring.h
index 74b2c1e4b..8c77c4499 100644
--- a/src/emscripten-optimizer/istring.h
+++ b/src/emscripten-optimizer/istring.h
@@ -137,14 +137,18 @@ struct IString {
   bool is() const     { return str != nullptr; }
   bool isNull() const { return str == nullptr; }
 
-  bool startsWith(const char *prefix) const {
+  const char* stripPrefix(const char *prefix) const {
     const char *ptr = str;
     while (true) {
-      if (*prefix == 0) return true;
-      if (*ptr == 0) return false;
-      if (*ptr++ != *prefix++) return false;
+      if (*prefix == 0) return ptr;
+      if (*ptr == 0) return nullptr;
+      if (*ptr++ != *prefix++) return nullptr;
     }
   }
+
+  bool startsWith(const char *prefix) const {
+    return stripPrefix(prefix) != nullptr;
+  }
 };
 
 } // namespace cashew
diff --git a/src/wasm-emscripten.cpp b/src/wasm-emscripten.cpp
index 0a898f7e4..eced376bb 100644
--- a/src/wasm-emscripten.cpp
+++ b/src/wasm-emscripten.cpp
@@ -29,6 +29,7 @@
 namespace wasm {
 
 cashew::IString EMSCRIPTEN_ASM_CONST("emscripten_asm_const");
+cashew::IString EM_JS_PREFIX("__em_js__");
 
 static constexpr const char* dummyFunction = "__wasm_nullptr";
 
@@ -296,6 +297,64 @@ void EmscriptenGlueGenerator::generateJSCallThunks(
       tableSegmentData.size();
 }
 
+std::vector<Address> getSegmentOffsets(Module& wasm) {
+  std::vector<Address> segmentOffsets;
+  for (unsigned i = 0; i < wasm.memory.segments.size(); ++i) {
+    Const* addrConst = wasm.memory.segments[i].offset->cast<Const>();
+    auto address = addrConst->value.geti32();
+    segmentOffsets.push_back(address);
+  }
+  return segmentOffsets;
+}
+
+std::string escape(const char *input) {
+  std::string code = input;
+  // replace newlines quotes with escaped newlines
+  size_t curr = 0;
+  while ((curr = code.find("\\n", curr)) != std::string::npos) {
+    code = code.replace(curr, 2, "\\\\n");
+    curr += 3; // skip this one
+  }
+  // replace double quotes with escaped single quotes
+  curr = 0;
+  while ((curr = code.find('"', curr)) != std::string::npos) {
+    if (curr == 0 || code[curr-1] != '\\') {
+      code = code.replace(curr, 1, "\\" "\"");
+      curr += 2; // skip this one
+    } else { // already escaped, escape the slash as well
+      code = code.replace(curr, 1, "\\" "\\" "\"");
+      curr += 3; // skip this one
+    }
+  }
+  return code;
+}
+
+const char* stringAtAddr(Module& wasm,
+                         std::vector<Address> const& segmentOffsets,
+                         Address address) {
+  for (unsigned i = 0; i < wasm.memory.segments.size(); ++i) {
+    Memory::Segment& segment = wasm.memory.segments[i];
+    Address offset = segmentOffsets[i];
+    if (address >= offset && address < offset + segment.data.size()) {
+      return &segment.data[address - offset];
+    }
+  }
+  return nullptr;
+}
+
+std::string codeForConstAddr(Module& wasm,
+                             std::vector<Address> const& segmentOffsets,
+                             Const* addrConst) {
+  auto address = addrConst->value.geti32();
+  const char* str = stringAtAddr(wasm, segmentOffsets, address);
+  if (!str) {
+    // If we can't find the segment corresponding with the address, then we
+    // omitted the segment and the address points to an empty string.
+    return escape("");
+  }
+  return escape(str);
+}
+
 struct AsmConstWalker : public PostWalker<AsmConstWalker> {
   Module& wasm;
   std::vector<Address> segmentOffsets; // segment index => address offset
@@ -304,31 +363,24 @@ struct AsmConstWalker : public PostWalker<AsmConstWalker> {
   std::map<std::string, Address> ids;
   std::set<std::string> allSigs;
 
-  AsmConstWalker(Module& _wasm) : wasm(_wasm) {
-    for (unsigned i = 0; i < wasm.memory.segments.size(); ++i) {
-      Const* addrConst = wasm.memory.segments[i].offset->cast<Const>();
-      auto address = addrConst->value.geti32();
-      segmentOffsets.push_back(address);
-    }
-  }
+  AsmConstWalker(Module& _wasm)
+    : wasm(_wasm),
+      segmentOffsets(getSegmentOffsets(wasm)) { }
 
   void visitCallImport(CallImport* curr);
 
 private:
-  std::string codeForConstAddr(Const* addrConst);
-  const char* stringAtAddr(Address adddress);
   Literal idLiteralForCode(std::string code);
   std::string asmConstSig(std::string baseSig);
   Name nameForImportWithSig(std::string sig);
   void addImport(Name importName, std::string baseSig);
-  std::string escape(const char *input);
 };
 
 void AsmConstWalker::visitCallImport(CallImport* curr) {
   Import* import = wasm.getImport(curr->target);
   if (import->base.hasSubstring(EMSCRIPTEN_ASM_CONST)) {
     auto arg = curr->operands[0]->cast<Const>();
-    auto code = codeForConstAddr(arg);
+    auto code = codeForConstAddr(wasm, segmentOffsets, arg);
     arg->value = idLiteralForCode(code);
     auto baseSig = getSig(curr);
     auto sig = asmConstSig(baseSig);
@@ -343,51 +395,6 @@ void AsmConstWalker::visitCallImport(CallImport* curr) {
   }
 }
 
-std::string AsmConstWalker::codeForConstAddr(Const* addrConst) {
-  auto address = addrConst->value.geti32();
-  const char* str = stringAtAddr(address);
-  if (!str) {
-    // If we can't find the segment corresponding with the address, then we
-    // omitted the segment and the address points to an empty string.
-    return escape("");
-  }
-  auto result = escape(str);
-  return result;
-}
-
-const char* AsmConstWalker::stringAtAddr(Address address) {
-  for (unsigned i = 0; i < wasm.memory.segments.size(); ++i) {
-    Memory::Segment &segment = wasm.memory.segments[i];
-    Address offset = segmentOffsets[i];
-    if (address >= offset && address < offset + segment.data.size()) {
-      return &segment.data[address - offset];
-    }
-  }
-  return nullptr;
-}
-
-std::string AsmConstWalker::escape(const char *input) {
-  std::string code = input;
-  // replace newlines quotes with escaped newlines
-  size_t curr = 0;
-  while ((curr = code.find("\\n", curr)) != std::string::npos) {
-    code = code.replace(curr, 2, "\\\\n");
-    curr += 3; // skip this one
-  }
-  // replace double quotes with escaped single quotes
-  curr = 0;
-  while ((curr = code.find('"', curr)) != std::string::npos) {
-    if (curr == 0 || code[curr-1] != '\\') {
-      code = code.replace(curr, 1, "\\" "\"");
-      curr += 2; // skip this one
-    } else { // already escaped, escape the slash as well
-      code = code.replace(curr, 1, "\\" "\\" "\"");
-      curr += 3; // skip this one
-    }
-  }
-  return code;
-}
-
 Literal AsmConstWalker::idLiteralForCode(std::string code) {
   int32_t id;
   if (ids.count(code) == 0) {
@@ -445,8 +452,60 @@ AsmConstWalker fixEmAsmConstsAndReturnWalker(Module& wasm) {
   return walker;
 }
 
+struct EmJsWalker : public PostWalker<EmJsWalker> {
+  Module& wasm;
+  std::vector<Address> segmentOffsets; // segment index => address offset
+
+  std::map<std::string, std::string> codeByName;
+
+  EmJsWalker(Module& _wasm)
+    : wasm(_wasm),
+      segmentOffsets(getSegmentOffsets(wasm)) { }
+
+  void visitFunction(Function* curr) {
+    if (!curr->name.startsWith(EM_JS_PREFIX.str)) {
+      return;
+    }
+    auto funcName = std::string(curr->name.stripPrefix(EM_JS_PREFIX.str));
+    auto addrConst = curr->body->dynCast<Const>();
+    if (addrConst == nullptr) {
+      auto block = curr->body->dynCast<Block>();
+      Expression* first = nullptr;
+      if (block && block->list.size() > 0) {
+        first = block->list[0];
+      }
+      if (first) {
+        addrConst = first->dynCast<Const>();
+      }
+    }
+    if (addrConst == nullptr) {
+      Fatal() << "Unexpected generated __em_js__ function body: " << curr;
+    }
+    auto code = codeForConstAddr(wasm, segmentOffsets, addrConst);
+    codeByName[funcName] = code;
+  }
+};
+
+EmJsWalker fixEmJsFuncsAndReturnWalker(Module& wasm) {
+  EmJsWalker walker(wasm);
+  walker.walkModule(&wasm);
+
+  std::vector<Name> toRemove;
+  for (auto& func : wasm.functions) {
+    if (func->name.startsWith(EM_JS_PREFIX.str)) {
+      toRemove.push_back(func->name);
+    }
+  }
+  for (auto funcName : toRemove) {
+    wasm.removeFunction(funcName);
+    wasm.removeExport(funcName);
+  }
+  return walker;
+}
+
 void EmscriptenGlueGenerator::fixEmAsmConsts() {
   fixEmAsmConstsAndReturnWalker(wasm);
+  fixEmJsFuncsAndReturnWalker(wasm);
 }
 
 template<class C>
@@ -492,6 +551,19 @@ std::string EmscriptenGlueGenerator::generateEmscriptenMetadata(
   }
   meta << "},";
 
+  EmJsWalker emJsWalker = fixEmJsFuncsAndReturnWalker(wasm);
+  if (emJsWalker.codeByName.size() > 0) {
+    meta << "\"emJsFuncs\": {";
+    commaFirst = true;
+    for (auto& pair : emJsWalker.codeByName) {
+      auto& name = pair.first;
+      auto& code = pair.second;
+      meta << maybeComma();
+      meta << '"' << name << "\": \"" << code << '"';
+    }
+    meta << "},";
+  }
+
   meta << "\"staticBump\": " << staticBump << ", ";
 
   meta << "\"initializers\": [";
@@ -519,8 +591,7 @@ std::string EmscriptenGlueGenerator::generateEmscriptenMetadata(
   commaFirst = true;
   for (const auto& import : wasm.imports) {
     if (import->kind == ExternalKind::Function &&
-        // TODO(jgravelle): uncomment when EM_JS is merged
-        // !import->name.startsWith(EM_JS_PREFIX.str) &&
+        (emJsWalker.codeByName.count(import->name.str) == 0) &&
         !import->name.startsWith(EMSCRIPTEN_ASM_CONST.str) &&
         !import->name.startsWith("invoke_") &&
         !import->name.startsWith("jsCall_")) {