From 3ae2e6ef5919198f5e0efc2aa86bb50b07eb2e72 Mon Sep 17 00:00:00 2001
From: Sam Clegg <sbc@chromium.org>
Date: Thu, 22 Jun 2023 20:37:08 -0700
Subject: [PATCH] Enable EM_JS in side modules

This works in a similar way to EM_ASM. See #18228.

Depends on https://github.com/WebAssembly/binaryen/pull/5780
---
 emcc.py                       |  1 +
 emscripten.py                 |  2 --
 src/library_dylink.js         | 27 ++++++++++++++++++++++++++-
 test/core/test_em_js_main.c   |  9 +++++++++
 test/core/test_em_js_main.out |  2 ++
 test/core/test_em_js_side.c   | 11 +++++++++++
 test/test_core.py             |  5 +++++
 tools/utils.py                |  7 +++++++
 8 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 test/core/test_em_js_main.c
 create mode 100644 test/core/test_em_js_main.out
 create mode 100644 test/core/test_em_js_side.c

diff --git a/emcc.py b/emcc.py
index 202732113a629..56710067c7941 100755
--- a/emcc.py
+++ b/emcc.py
@@ -831,6 +831,7 @@ def process_dynamic_libs(dylibs, lib_dirs):
   for dylib in dylibs:
     exports = webassembly.get_exports(dylib)
     exports = set(e.name for e in exports)
+    exports = [utils.removeprefix(e, '__em_js__') for e in exports]
     settings.SIDE_MODULE_EXPORTS.extend(sorted(exports))
 
     imports = webassembly.get_imports(dylib)
diff --git a/emscripten.py b/emscripten.py
index 43d9d51c68ec6..6e22c5ce09642 100644
--- a/emscripten.py
+++ b/emscripten.py
@@ -344,8 +344,6 @@ def emscript(in_wasm, out_wasm, outfile_js, memfile, js_syms):
           diagnostics.warning('em-js-i64', 'using 64-bit arguments in EM_JS function without WASM_BIGINT is not yet fully supported: `%s` (%s, %s)', em_js_func, c_sig, signature.params)
 
   if settings.SIDE_MODULE:
-    if metadata.emJsFuncs:
-      exit_with_error('EM_JS is not supported in side modules')
     logger.debug('emscript: skipping remaining js glue generation')
     return
 
diff --git a/src/library_dylink.js b/src/library_dylink.js
index 043328302ad00..a6ad95d3eadeb 100644
--- a/src/library_dylink.js
+++ b/src/library_dylink.js
@@ -178,7 +178,9 @@ var LibraryDylink = {
       '__wasm_call_ctors',
       '__start_em_asm',
       '__stop_em_asm',
-    ].includes(symName)
+      '__start_em_js',
+      '__stop_em_js',
+    ].includes(symName) || symName.startsWith('__em_js__')
 #if SPLIT_MODULE
         // Exports synthesized by wasm-split should be prefixed with '%'
         || symName[0] == '%'
@@ -779,6 +781,29 @@ var LibraryDylink = {
             start = HEAPU8.indexOf(0, start) + 1;
           }
         }
+
+        function addEmJs(name, sig, body) {
+          sig = sig.slice(1, -1).split(',');
+          var args = [];
+          for (var i in sig) {
+            args.push(sig[i].split(' ').pop());
+          }
+          var func = `(${args}) => ${body};`;
+#if DYLINK_DEBUG
+          dbg(`adding new EM_JS function: ${name} = ${func}`);
+#endif
+          {{{ makeEval('wasmImports[name] = eval(func)') }}};
+        }
+
+        for (var name in moduleExports) {
+          if (name.startsWith('__em_js__')) {
+            var start = moduleExports[name]
+            {{{ from64('start') }}}
+            var jsString = UTF8ToString(start);
+            var parts = jsString.split('<::>');
+            addEmJs(name.replace('__em_js__', ''), parts[0], parts[1]);
+          }
+        }
 #endif
 
         // initialize the module
diff --git a/test/core/test_em_js_main.c b/test/core/test_em_js_main.c
new file mode 100644
index 0000000000000..71048476e7555
--- /dev/null
+++ b/test/core/test_em_js_main.c
@@ -0,0 +1,9 @@
+#include <stdio.h>
+#include <emscripten.h>
+
+int test_side();
+
+int main() {
+  printf("in main\n");
+  return test_side();
+}
diff --git a/test/core/test_em_js_main.out b/test/core/test_em_js_main.out
new file mode 100644
index 0000000000000..bbe0ca2d46f9c
--- /dev/null
+++ b/test/core/test_em_js_main.out
@@ -0,0 +1,2 @@
+in main
+hello from side module 42 + hello
diff --git a/test/core/test_em_js_side.c b/test/core/test_em_js_side.c
new file mode 100644
index 0000000000000..3cc5ace0d3845
--- /dev/null
+++ b/test/core/test_em_js_side.c
@@ -0,0 +1,11 @@
+#include <emscripten/em_js.h>
+#include <stdio.h>
+
+EM_JS(void*, js_side_func, (int num, char* ptr), {
+  out(`hello from side module ${num} + ${UTF8ToString(ptr)}`);
+  return 99;
+});
+
+void test_side() {
+  js_side_func(42, "hello");
+}
diff --git a/test/test_core.py b/test/test_core.py
index 851356f0333e9..fd831793ba9e0 100644
--- a/test/test_core.py
+++ b/test/test_core.py
@@ -2322,6 +2322,11 @@ def test_em_js_address_taken(self):
       self.set_setting('MAIN_MODULE', 2)
       self.do_core_test('test_em_js_address_taken.c')
 
+  @needs_dylink
+  def test_em_js_side_module(self):
+    self.build(test_file('core/test_em_js_side.c'), js_outfile=False, emcc_args=['-sSIDE_MODULE'], output_basename='side')
+    self.do_core_test('test_em_js_main.c', emcc_args=['-sMAIN_MODULE=2', 'side.wasm'])
+
   def test_runtime_stacksave(self):
     self.do_runf(test_file('core/test_runtime_stacksave.c'), 'success')
 
diff --git a/tools/utils.py b/tools/utils.py
index bd688cc5c1157..f3b8294195237 100644
--- a/tools/utils.py
+++ b/tools/utils.py
@@ -29,6 +29,13 @@ def safe_ensure_dirs(dirname):
   os.makedirs(dirname, exist_ok=True)
 
 
+# TODO(sbc): Replace with str.removeprefix once we update to python3.9
+def removeprefix(string, prefix):
+  if string.startswith(prefix):
+    return string[len(prefix):]
+  return string
+
+
 @contextlib.contextmanager
 def chdir(dir):
   """A context manager that performs actions in the given directory."""