Support ASYNCIFY_IMPORTS and add ASYNCIFY_EXPORTS for JSPI. (#18165)

Pass asyncified imports and exports on to the JSPI wasm-opt pass so
only the specified functions are wrapped.

* Update src/library_async.js

Co-authored-by: Alon Zakai <[email protected]>

Author: brendandahl

emcc.py

@@ -105,6 +105,10 @@
   '_dlopen_js', '__asyncjs__*'
 ]
 
+DEFAULT_ASYNCIFY_EXPORTS = [
+  'main'
+]
+
 # Target options
 final_js = None
 
@@ -665,6 +669,8 @@ def check_human_readable_list(items):
       passes += ['--pass-arg=asyncify-onlylist@%s' % ','.join(settings.ASYNCIFY_ONLY)]
   elif settings.ASYNCIFY == 2:
     passes += ['--jspi']
+    passes += ['--pass-arg=jspi-imports@%s' % ','.join(settings.ASYNCIFY_IMPORTS)]
+    passes += ['--pass-arg=jspi-exports@%s' % ','.join(settings.ASYNCIFY_EXPORTS)]
   if settings.BINARYEN_IGNORE_IMPLICIT_TRAPS:
     passes += ['--ignore-implicit-traps']
   # normally we can assume the memory, if imported, has not been modified
@@ -2650,6 +2656,8 @@ def check_memory_setting(setting):
       settings.ASYNCIFY_IMPORTS += ['invoke_*']
     # add the default imports
     settings.ASYNCIFY_IMPORTS += DEFAULT_ASYNCIFY_IMPORTS
+    # add the default exports (only used for ASYNCIFY == 2)
+    settings.ASYNCIFY_EXPORTS += DEFAULT_ASYNCIFY_EXPORTS
 
     # return the full import name, including module. The name may
     # already have a module prefix; if not, we assume it is "env".

src/library_async.js

@@ -43,23 +43,24 @@ mergeInto(LibraryManager.library, {
             var isAsyncifyImport = ASYNCIFY_IMPORTS.indexOf(x) >= 0 ||
                                    x.startsWith('__asyncjs__');
 #if ASYNCIFY == 2
-            // Wrap all imports with a suspending WebAssembly function.
-            // TODO: wrap only async functions.
+            // Wrap async imports with a suspending WebAssembly function.
+            if (isAsyncifyImport) {
 #if ASSERTIONS
-            assert(sig, 'Missing __sig for ' + x);
+              assert(sig, 'Missing __sig for ' + x);
 #endif
-            var type = sigToWasmTypes(sig);
+              var type = sigToWasmTypes(sig);
 #if ASYNCIFY_DEBUG
-            dbg('asyncify: suspendOnReturnedPromise for', x, original);
-#endif
-            // Add space for the suspender promise that will be used in the
-            // Wasm wrapper function.
-            type.parameters.unshift('externref');
-            imports[x] = original = new WebAssembly.Function(
-              type,
-              original,
-              { suspending: 'first' }
-            );
+              dbg('asyncify: suspendOnReturnedPromise for', x, original);
+#endif
+              // Add space for the suspender promise that will be used in the
+              // Wasm wrapper function.
+              type.parameters.unshift('externref');
+              imports[x] = original = new WebAssembly.Function(
+                type,
+                original,
+                { suspending: 'first' }
+              );
+            }
 #endif
 #if ASSERTIONS && ASYNCIFY != 2 // We cannot apply assertions with stack switching, as the imports must not be modified from suspender.suspendOnReturnedPromise TODO find a way
             imports[x] = function() {
@@ -102,6 +103,9 @@ mergeInto(LibraryManager.library, {
     instrumentWasmExports: function(exports) {
 #if ASYNCIFY_DEBUG
       dbg('asyncify instrumenting exports');
+#endif
+#if ASYNCIFY == 2
+      var ASYNCIFY_EXPORTS = {{{ JSON.stringify(ASYNCIFY_EXPORTS) }}};
 #endif
       var ret = {};
       for (var x in exports) {
@@ -110,22 +114,25 @@ mergeInto(LibraryManager.library, {
           if (typeof original == 'function') {
 #if ASYNCIFY == 2
             // Wrap all exports with a promising WebAssembly function.
-            // TODO: wrap only async functions.
+            var isAsyncifyExport = ASYNCIFY_EXPORTS.indexOf(x) >= 0;
+            if (isAsyncifyExport) {
 #if ASYNCIFY_DEBUG
-            dbg('asyncify: returnPromiseOnSuspend for', x, original);
+              dbg('asyncify: returnPromiseOnSuspend for', x, original);
 #endif
-            var type = WebAssembly.Function.type(original);
-            var parameters = type.parameters;
-            var results = type.results;
+              var type = WebAssembly.Function.type(original);
+              var parameters = type.parameters;
+              var results = type.results;
 #if ASSERTIONS
-            assert(results.length !== 0, 'There must be a return result')
-#endif
-            // Remove the extern ref;
-            parameters.shift();
-            original = new WebAssembly.Function(
-              { parameters , results: ['externref'] },
-              original,
-              { promising : 'first' });
+              assert(results.length !== 0, 'There must be a return result')
+              assert(parameters[0] === 'externref', 'First param must be externref.');
+#endif
+              // Remove the extern ref.
+              parameters.shift();
+              original = new WebAssembly.Function(
+                { parameters , results: ['externref'] },
+                original,
+                { promising : 'first' });
+            }
 #endif
             ret[x] = function() {
 #if ASYNCIFY_DEBUG >= 2

src/settings.js

@@ -864,6 +864,12 @@ var ASYNCIFY_LAZY_LOAD_CODE = false;
 // [link]
 var ASYNCIFY_DEBUG = 0;
 
+// Specify which of the exports will have JSPI applied to them and return a
+// promise.
+// Only supported for ASYNCIFY==2 mode.
+// [link]
+var ASYNCIFY_EXPORTS = [];
+
 // Runtime elements that are exported on Module by default. We used to export
 // quite a lot here, but have removed them all. You should use
 // EXPORTED_RUNTIME_METHODS for things you want to export from the runtime.