3 files changed, 55 insertions, 11 deletions

diff --git a/src/passes/LegalizeJSInterface.cpp b/src/passes/LegalizeJSInterface.cpp
index e5a69beb1..dd702ba7d 100644
--- a/src/passes/LegalizeJSInterface.cpp
+++ b/src/passes/LegalizeJSInterface.cpp
@@ -43,16 +43,32 @@
 
 namespace wasm {
 
+// These are aliases for getTempRet0/setTempRet0 which emscripten defines in
+// compiler-rt and exports under these names.
+static Name GET_TEMP_RET_EXPORT("__get_temp_ret");
+static Name SET_TEMP_RET_EXPORT("__set_temp_ret");
+
+// For non-emscripten module we expect the host to define these functions so
+// and we import them under these names.
+static Name GET_TEMP_RET_IMPORT("getTempRet0");
+static Name SET_TEMP_RET_IMPORT("setTempRet0");
+
 struct LegalizeJSInterface : public Pass {
   bool full;
 
   LegalizeJSInterface(bool full) : full(full) {}
 
   void run(PassRunner* runner, Module* module) override {
+    setTempRet0 = nullptr;
+    getTempRet0 = nullptr;
     auto exportOriginals =
       !runner->options
          .getArgumentOrDefault("legalize-js-interface-export-originals", "")
          .empty();
+    exportedHelpers =
+      !runner->options
+         .getArgumentOrDefault("legalize-js-interface-exported-helpers", "")
+         .empty();
     // for each illegal export, we must export a legalized stub instead
     std::vector<std::unique_ptr<Export>> newExports;
     for (auto& ex : module->exports) {
@@ -82,6 +98,7 @@ struct LegalizeJSInterface : public Pass {
         }
       }
     }
+
     for (auto& ex : newExports) {
       module->addExport(std::move(ex));
     }
@@ -150,11 +167,17 @@ struct LegalizeJSInterface : public Pass {
         module->removeFunction(pair.first);
       }
     }
+
+    module->removeExport(GET_TEMP_RET_EXPORT);
+    module->removeExport(SET_TEMP_RET_EXPORT);
   }
 
 private:
   // map of illegal to legal names for imports
   std::map<Name, Name> illegalImportsToLegal;
+  bool exportedHelpers = false;
+  Function* getTempRet0 = nullptr;
+  Function* setTempRet0 = nullptr;
 
   template<typename T> bool isIllegal(T* t) {
     for (const auto& param : t->getParams()) {
@@ -185,6 +208,32 @@ private:
     return im->module == ENV && im->base.startsWith("invoke_");
   }
 
+  Function* tempSetter(Module* module) {
+    if (!setTempRet0) {
+      if (exportedHelpers) {
+        auto* ex = module->getExport(SET_TEMP_RET_EXPORT);
+        setTempRet0 = module->getFunction(ex->value);
+      } else {
+        setTempRet0 = getFunctionOrImport(
+          module, SET_TEMP_RET_IMPORT, Type::i32, Type::none);
+      }
+    }
+    return setTempRet0;
+  }
+
+  Function* tempGetter(Module* module) {
+    if (!getTempRet0) {
+      if (exportedHelpers) {
+        auto* ex = module->getExport(GET_TEMP_RET_EXPORT);
+        getTempRet0 = module->getFunction(ex->value);
+      } else {
+        getTempRet0 = getFunctionOrImport(
+          module, GET_TEMP_RET_IMPORT, Type::none, Type::i32);
+      }
+    }
+    return getTempRet0;
+  }
+
   // JS calls the export, so it must call a legal stub that calls the actual
   // wasm function
   Name makeLegalStub(Function* func, Module* module) {
@@ -220,13 +269,13 @@ private:
       func->getResults() == Type::i64 ? Type::i32 : func->getResults();
     legal->type = Signature(Type(legalParams), resultsType);
     if (func->getResults() == Type::i64) {
-      Function* f =
-        getFunctionOrImport(module, SET_TEMP_RET0, Type::i32, Type::none);
       auto index = Builder::addVar(legal, Name(), Type::i64);
       auto* block = builder.makeBlock();
       block->list.push_back(builder.makeLocalSet(index, call));
-      block->list.push_back(builder.makeCall(
-        f->name, {I64Utilities::getI64High(builder, index)}, Type::none));
+      block->list.push_back(
+        builder.makeCall(tempSetter(module)->name,
+                         {I64Utilities::getI64High(builder, index)},
+                         Type::none));
       block->list.push_back(I64Utilities::getI64Low(builder, index));
       block->finalize();
       legal->body = block;
@@ -267,10 +316,9 @@ private:
     }
 
     if (im->getResults() == Type::i64) {
-      Function* f =
-        getFunctionOrImport(module, GET_TEMP_RET0, Type::none, Type::i32);
       call->type = Type::i32;
-      Expression* get = builder.makeCall(f->name, {}, call->type);
+      Expression* get =
+        builder.makeCall(tempGetter(module)->name, {}, call->type);
       stub->body = I64Utilities::recreateI64(builder, call, get);
     } else {
       call->type = im->getResults();
diff --git a/src/shared-constants.h b/src/shared-constants.h
index 6ad3ab9ac..5b85dd74d 100644
--- a/src/shared-constants.h
+++ b/src/shared-constants.h
@@ -24,8 +24,6 @@ namespace wasm {
 extern Name MEMORY_BASE;
 extern Name TABLE_BASE;
 extern Name STACK_POINTER;
-extern Name GET_TEMP_RET0;
-extern Name SET_TEMP_RET0;
 extern Name NEW_SIZE;
 extern Name MODULE;
 extern Name START;
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index 34b1e82b7..1c9c1389d 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -57,8 +57,6 @@ const char* StringsFeature = "strings";
 Name MEMORY_BASE("__memory_base");
 Name TABLE_BASE("__table_base");
 Name STACK_POINTER("__stack_pointer");
-Name GET_TEMP_RET0("getTempRet0");
-Name SET_TEMP_RET0("setTempRet0");
 Name NEW_SIZE("newSize");
 Name MODULE("module");
 Name START("start");