LegalizeJSInterface: Look for get/setTempRet0 as exports (#4881)

This allows emscripten to move these helper functions from JS library
imports to native wasm exports.

See https://github.com/emscripten-core/emscripten/issues/7273

4 files changed, 128 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");
diff --git a/test/lit/passes/legalize-js-interface-exported-helpers.wast b/test/lit/passes/legalize-js-interface-exported-helpers.wast
new file mode 100644
index 000000000..253c3a5a7
--- /dev/null
+++ b/test/lit/passes/legalize-js-interface-exported-helpers.wast
@@ -0,0 +1,73 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_test.py and could be cleaned up.
+
+;; Check that legalize-js-interface-exported-helpers ends up using the
+;; existing __set_temp_ret/__get_temp_ret helpers.
+;; RUN: wasm-opt %s --legalize-js-interface --pass-arg=legalize-js-interface-exported-helpers -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $none_=>_i32 (func (result i32)))
+
+  ;; CHECK:      (type $none_=>_i64 (func (result i64)))
+
+  ;; CHECK:      (type $i32_=>_none (func (param i32)))
+
+  ;; CHECK:      (import "env" "imported" (func $legalimport$imported (result i32)))
+
+  ;; CHECK:      (export "get_i64" (func $legalstub$get_i64))
+  (export "get_i64" (func $get_i64))
+  (import "env" "imported" (func $imported (result i64)))
+  (export "__set_temp_ret" (func $__set_temp_ret))
+  (export "__get_temp_ret" (func $__get_temp_ret))
+  ;; CHECK:      (func $get_i64 (result i64)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $legalfunc$imported)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i64.const 0)
+  ;; CHECK-NEXT: )
+  (func $get_i64 (result i64)
+    (drop (call $imported))
+    (i64.const 0)
+  )
+  ;; CHECK:      (func $__set_temp_ret (param $0 i32)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $__set_temp_ret (param i32))
+  ;; CHECK:      (func $__get_temp_ret (result i32)
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $__get_temp_ret (result i32)
+    (i32.const 0)
+  )
+)
+;; CHECK:      (func $legalstub$get_i64 (result i32)
+;; CHECK-NEXT:  (local $0 i64)
+;; CHECK-NEXT:  (local.set $0
+;; CHECK-NEXT:   (call $get_i64)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (call $__set_temp_ret
+;; CHECK-NEXT:   (i32.wrap_i64
+;; CHECK-NEXT:    (i64.shr_u
+;; CHECK-NEXT:     (local.get $0)
+;; CHECK-NEXT:     (i64.const 32)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT:  (i32.wrap_i64
+;; CHECK-NEXT:   (local.get $0)
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )
+
+;; CHECK:      (func $legalfunc$imported (result i64)
+;; CHECK-NEXT:  (i64.or
+;; CHECK-NEXT:   (i64.extend_i32_u
+;; CHECK-NEXT:    (call $legalimport$imported)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:   (i64.shl
+;; CHECK-NEXT:    (i64.extend_i32_u
+;; CHECK-NEXT:     (call $__get_temp_ret)
+;; CHECK-NEXT:    )
+;; CHECK-NEXT:    (i64.const 32)
+;; CHECK-NEXT:   )
+;; CHECK-NEXT:  )
+;; CHECK-NEXT: )