Unify module splitting implementation for JSPI (#26396)

Previously, module splitting with JSPI required a special built-in
function `__load_secondary_module` and wasm-split setting to coordinate
the loading of the secondary module.

This change updates the `splitModuleProxyHandler` to support JSPI by
wrapping the generated placeholder functions in
`WebAssembly.Suspending`. The placeholder will then automatically load
the secondary module asynchronously.

A standalone `save_profile_data.js` script is extracted from existing
tests for reuse, and a new test `test_split_module_embind_jspi` is added
to ensure module splitting functions correctly alongside embind and
JSPI.

---------

Co-authored-by: Heejin Ahn <aheejin@gmail.com>

Author: brendandahl

src/lib/libasync.js

@@ -514,17 +514,6 @@ addToLibrary({
     });
   },
 
-  _load_secondary_module__sig: 'v',
-  _load_secondary_module__async: 'auto',
-  _load_secondary_module: async function() {
-    // Mark the module as loading for the wasm module (so it doesn't try to load it again).
-    wasmExports['load_secondary_module_status'].value = 1;
-    var imports = {'primary': wasmRawExports};
-    // Replace '.wasm' suffix with '.deferred.wasm'.
-    var deferred = wasmBinaryFile.slice(0, -5) + '.deferred.wasm';
-    await instantiateAsync(null, deferred, imports);
-  },
-
   $Fibers__deps: ['$Asyncify', 'emscripten_stack_set_limits', '$stackRestore'],
   $Fibers: {
     nextFiber: 0,

src/preamble.js

@@ -506,7 +506,7 @@ async function getWasmBinary(binaryFile) {
 #endif
 
 #if SPLIT_MODULE
-{{{ makeModuleReceiveWithVar('loadSplitModule', undefined, 'instantiateSync') }}}
+{{{ makeModuleReceiveWithVar('loadSplitModule', undefined, JSPI ? '(secondaryFile, imports) => instantiateAsync(null, secondaryFile, imports)' : 'instantiateSync') }}}
 var splitModuleProxyHandler = {
   get(target, moduleName, receiver) {
     if (moduleName.startsWith('placeholder')) {
@@ -519,22 +519,23 @@ var splitModuleProxyHandler = {
       }
       return new Proxy({}, {
         get(target, base, receiver) {
-          return (...args) => {
-#if ASYNCIFY == 2
-            throw new Error('Placeholder function "' + base + '" should not be called when using JSPI.');
-#else
+          let ret = {{{ asyncIf(ASYNCIFY == 2) }}} (...args) => {
 #if RUNTIME_DEBUG
             dbg(`placeholder function called: ${base}`);
 #endif
             var imports = {'primary': wasmRawExports};
             // Replace '.wasm' suffix with '.deferred.wasm'.
-            loadSplitModule(secondaryFile, imports, base);
+            {{{ awaitIf(ASYNCIFY == 2) }}}loadSplitModule(secondaryFile, imports, base);
 #if RUNTIME_DEBUG
             dbg('instantiated deferred module, continuing');
 #endif
             return wasmTable.get({{{ toIndexType('base') }}})(...args);
+          };
+#if JSPI
+          return new WebAssembly.Suspending(ret);
+#else
+          return ret;
 #endif
-          }
         }
       });
     }

test/other/save_profile_data.js

@@ -0,0 +1,18 @@
+function save_profile_data() {
+  var __write_profile = wasmExports['__write_profile'];
+  if (__write_profile) {
+    var len = __write_profile({{{ to64('0') }}}, 0);
+    var offset = _malloc(len);
+    var actualLen = __write_profile({{{ to64('offset') }}}, len);
+    var profile_data = HEAPU8.subarray(offset, offset + len);
+    if (typeof writeFile !== 'undefined') {
+      writeFile('profile.data', profile_data);
+    } else if (typeof fs !== 'undefined') {
+      fs.writeFileSync('profile.data', profile_data);
+    } else {
+      console.log(JSON.stringify(Array.from(profile_data)));
+    }
+    console.log('wrote profile data')
+    _free(offset);
+  }
+}

test/other/test_split_module.post.js

@@ -1,25 +1,8 @@
 #preprocess
+#include "save_profile_data.js"
 
 function saveProfileData() {
-  var __write_profile = wasmExports['__write_profile'];
-  if (__write_profile) {
-    var len = __write_profile({{{ to64('0') }}}, 0);
-    var offset = _malloc(len);
-    var actualLen = __write_profile({{{ to64('offset') }}}, len);
-    var profile_data = HEAPU8.subarray(offset, offset + len);
-    if (typeof writeFile !== 'undefined') {
-      console.log('using writeFile')
-      writeFile('profile.data', profile_data);
-    } else if (typeof fs !== 'undefined') {
-      console.log('using fs.writeFileSync')
-      fs.writeFileSync('profile.data', profile_data);
-    } else {
-      console.log(JSON.stringify(Array.from(profile_data)));
-    }
-    console.log('profile size is', actualLen, 'bytes (allocated', len, 'bytes)');
-    console.log('wrote profile data')
-    _free(offset);
-  }
+  save_profile_data();
 
   // Say hello *after* recording the profile so that all functions are deferred.
   var result = _say_hello();

test/other/test_split_module_embind_jspi.cpp

@@ -0,0 +1,15 @@
+#include <emscripten.h>
+#include <emscripten/bind.h>
+
+int primary_function() {
+  return 42;
+}
+
+int deferred_function() {
+  return 82;
+}
+
+EMSCRIPTEN_BINDINGS(module_splitting) {
+  emscripten::function("primary_function", &primary_function);
+  emscripten::function("deferred_function", &deferred_function, emscripten::async());
+}

test/other/test_split_module_embind_jspi.post.js

@@ -0,0 +1,13 @@
+#preprocess
+#include "save_profile_data.js"
+
+async function saveProfileData() {
+  console.log('primary_function: ' + Module.primary_function());
+  save_profile_data();
+  // deferred_function *after* recording the profile so that all functions are deferred.
+  var result = Module.deferred_function();
+  console.log('deferred_function: ' + result);
+  console.log('deferred_function await: ' + await result);
+}
+
+addOnPostRun(saveProfileData);

test/other/test_split_module_embind_jspi.pre.js

@@ -0,0 +1,5 @@
+Module["loadSplitModule"] = async function(deferred, imports, prop) {
+    console.log('Custom handler for loading split module.');
+
+    return instantiateAsync(null, deferred, imports);
+}

test/test_other.py

@@ -12376,8 +12376,6 @@ def test_split_module(self, customLoader, jspi, opt):
     wasm_split_run = [wasm_split, '-g',
                       '--enable-mutable-globals', '--enable-bulk-memory', '--enable-nontrapping-float-to-int',
                       '--export-prefix=%', 'test_split_module.wasm.orig', '-o1', 'primary.wasm', '-o2', 'secondary.wasm', '--profile=profile.data']
-    if jspi:
-      wasm_split_run += ['--jspi', '--enable-reference-types']
     if self.get_setting('MEMORY64'):
       wasm_split_run += ['--enable-memory64']
     self.run_process(wasm_split_run)
@@ -12435,6 +12433,38 @@ def test_split_main_module(self):
     self.assertIn('Hello from main!', result)
     self.assertIn('Hello from lib!', result)
 
+  @also_with_wasm64
+  @requires_jspi
+  def test_split_module_embind_jspi(self):
+    self.set_setting('SPLIT_MODULE')
+    self.cflags += ['-Wno-experimental']
+    self.cflags += ['--post-js', test_file('other/test_split_module_embind_jspi.post.js')]
+    self.cflags += ['--pre-js', test_file('other/test_split_module_embind_jspi.pre.js')]
+    self.cflags += ['-sEXPORTED_FUNCTIONS=_malloc,_free']
+    self.cflags += ['-lembind']
+    expected_pre_split_output = 'primary_function: 42\n'
+    expected_post_split_output = ('deferred_function: [object Promise]\n'
+                                  'deferred_function await: 82\n')
+    expected_output = (expected_pre_split_output +
+                       'wrote profile data\n' +
+                       expected_post_split_output)
+    result = self.do_runf('other/test_split_module_embind_jspi.cpp', expected_output=expected_output)
+    self.assertExists('test_split_module_embind_jspi.wasm')
+    self.assertExists('test_split_module_embind_jspi.wasm.orig')
+    self.assertExists('profile.data')
+
+    wasm_split = os.path.join(building.get_binaryen_bin(), 'wasm-split')
+    wasm_split_run = [wasm_split, '-g',
+                      '--enable-mutable-globals', '--enable-bulk-memory', '--enable-nontrapping-float-to-int',
+                      '--export-prefix=%', 'test_split_module_embind_jspi.wasm.orig', '-o1', 'test_split_module_embind_jspi.wasm', '-o2', 'test_split_module_embind_jspi.deferred.wasm', '--profile=profile.data']
+    self.run_process(wasm_split_run)
+    result = self.run_js('test_split_module_embind_jspi.js')
+    self.assertNotIn('profile', result)
+    self.assertIn((expected_pre_split_output +
+                   'Custom handler for loading split module.\n' +
+                   expected_post_split_output),
+                  result)
+
   @crossplatform
   @flaky('https://github.com/emscripten-core/emscripten/issues/25206')
   def test_gen_struct_info(self):

tools/emscripten.py

@@ -842,11 +842,6 @@ def add_standard_wasm_imports(send_items_map):
       'memory_grow_post',
     ]
 
-  if settings.SPLIT_MODULE and settings.ASYNCIFY == 2:
-    # Calls to this function are generated by binaryen so it must be manually
-    # imported.
-    extra_sent_items.append('__load_secondary_module')
-
   for s in extra_sent_items:
     send_items_map[s] = s
 

tools/link.py

@@ -1635,8 +1635,6 @@ def limit_incoming_module_api():
 
   if settings.SPLIT_MODULE:
     settings.INCOMING_MODULE_JS_API += ['loadSplitModule']
-    if settings.ASYNCIFY == 2:
-      settings.DEFAULT_LIBRARY_FUNCS_TO_INCLUDE += ['_load_secondary_module']
 
   # wasm side modules have suffix .wasm
   if settings.SIDE_MODULE and utils.suffix(target) in ('.js', '.mjs'):