Allow async main entry point with JavaScriptEventLoop

This change allows the Swift program to have an async main entry point
when the JavaScriptEventLoop is installed as the global executor.

Author: kateinoigakukun

Runtime/src/index.ts

@@ -45,6 +45,27 @@ export class SwiftRuntime {
         }
     }
 
+    main() {
+        const instance = this.instance;
+        try {
+            if (typeof instance.exports.main === "function") {
+                instance.exports.main();
+            } else if (
+                typeof instance.exports.__main_argc_argv === "function"
+            ) {
+                // Swift 6.0 and later use `__main_argc_argv` instead of `main`.
+                instance.exports.__main_argc_argv(0, 0);
+            }
+        } catch (error) {
+            if (error instanceof UnsafeEventLoopYield) {
+                // Ignore the error
+                return;
+            }
+            // Rethrow other errors
+            throw error;
+        }
+    }
+
     private get instance() {
         if (!this._instance)
             throw new Error("WebAssembly instance is not set yet");
@@ -419,5 +440,23 @@ export class SwiftRuntime {
                 signed ? BigInt.asIntN(64, value) : BigInt.asUintN(64, value)
             );
         },
+        swjs_unsafe_event_loop_yield: () => {
+            throw new UnsafeEventLoopYield();
+        },
     };
 }
+
+/// This error is thrown when yielding event loop control from `swift_task_asyncMainDrainQueue`
+/// to JavaScript. This is usually thrown when:
+/// - The entry point of the Swift program is `func main() async`
+/// - The Swift Concurrency's global executor is hooked by `JavaScriptEventLoop.installGlobalExecutor()`
+/// - Calling exported `main` or `__main_argc_argv` function from JavaScript
+///
+/// This exception must be caught by the caller of the exported function and the caller should
+/// catch this exception and just ignore it.
+///
+/// FAQ: Why this error is thrown?
+/// This error is thrown to unwind the call stack of the Swift program and return the control to
+/// the JavaScript side. Otherwise, the `swift_task_asyncMainDrainQueue` ends up with `abort()`
+/// because the event loop expects `exit()` call before the end of the event loop.
+class UnsafeEventLoopYield extends Error {}

Runtime/src/js-value.ts

@@ -82,7 +82,13 @@ export const write = (
     is_exception: boolean,
     memory: Memory
 ) => {
-    const kind = writeAndReturnKindBits(value, payload1_ptr, payload2_ptr, is_exception, memory);
+    const kind = writeAndReturnKindBits(
+        value,
+        payload1_ptr,
+        payload2_ptr,
+        is_exception,
+        memory
+    );
     memory.writeUint32(kind_ptr, kind);
 };
 

Runtime/src/types.ts

@@ -102,6 +102,7 @@ export interface ImportedFunctions {
     swjs_i64_to_bigint(value: bigint, signed: bool): ref;
     swjs_bigint_to_i64(ref: ref, signed: bool): bigint;
     swjs_i64_to_bigint_slow(lower: number, upper: number, signed: bool): ref;
+    swjs_unsafe_event_loop_yield: () => void;
 }
 
 export const enum LibraryFeatures {

Sources/JavaScriptEventLoop/JavaScriptEventLoop.swift

@@ -1,5 +1,6 @@
 import JavaScriptKit
 import _CJavaScriptEventLoop
+import _CJavaScriptKit
 
 // NOTE: `@available` annotations are semantically wrong, but they make it easier to develop applications targeting WebAssembly in Xcode.
 
@@ -90,6 +91,14 @@ public final class JavaScriptEventLoop: SerialExecutor, @unchecked Sendable {
     public static func installGlobalExecutor() {
         guard !didInstallGlobalExecutor else { return }
 
+        #if compiler(>=5.9)
+        typealias swift_task_asyncMainDrainQueue_hook_Fn = @convention(thin) (swift_task_asyncMainDrainQueue_original, swift_task_asyncMainDrainQueue_override) -> Void
+        let swift_task_asyncMainDrainQueue_hook_impl: swift_task_asyncMainDrainQueue_hook_Fn = { _, _ in
+            _unsafe_event_loop_yield()
+        }
+        swift_task_asyncMainDrainQueue_hook = unsafeBitCast(swift_task_asyncMainDrainQueue_hook_impl, to: UnsafeMutableRawPointer?.self)
+        #endif
+
         typealias swift_task_enqueueGlobal_hook_Fn = @convention(thin) (UnownedJob, swift_task_enqueueGlobal_original) -> Void
         let swift_task_enqueueGlobal_hook_impl: swift_task_enqueueGlobal_hook_Fn = { job, original in
             JavaScriptEventLoop.shared.unsafeEnqueue(job)

Sources/JavaScriptKit/Runtime/index.js

@@ -27,13 +27,13 @@ typedef SWIFT_CC(swift) void (*swift_task_enqueueGlobal_original)(
     Job *_Nonnull job);
 
 SWIFT_EXPORT_FROM(swift_Concurrency)
-void *_Nullable swift_task_enqueueGlobal_hook;
+extern void *_Nullable swift_task_enqueueGlobal_hook;
 
 /// A hook to take over global enqueuing with delay.
 typedef SWIFT_CC(swift) void (*swift_task_enqueueGlobalWithDelay_original)(
     unsigned long long delay, Job *_Nonnull job);
 SWIFT_EXPORT_FROM(swift_Concurrency)
-void *_Nullable swift_task_enqueueGlobalWithDelay_hook;
+extern void *_Nullable swift_task_enqueueGlobalWithDelay_hook;
 
 typedef SWIFT_CC(swift) void (*swift_task_enqueueGlobalWithDeadline_original)(
     long long sec,
@@ -42,12 +42,23 @@ typedef SWIFT_CC(swift) void (*swift_task_enqueueGlobalWithDeadline_original)(
     long long tnsec,
     int clock, Job *_Nonnull job);
 SWIFT_EXPORT_FROM(swift_Concurrency)
-void *_Nullable swift_task_enqueueGlobalWithDeadline_hook;
+extern void *_Nullable swift_task_enqueueGlobalWithDeadline_hook;
 
 /// A hook to take over main executor enqueueing.
 typedef SWIFT_CC(swift) void (*swift_task_enqueueMainExecutor_original)(
     Job *_Nonnull job);
 SWIFT_EXPORT_FROM(swift_Concurrency)
-void *_Nullable swift_task_enqueueMainExecutor_hook;
+extern void *_Nullable swift_task_enqueueMainExecutor_hook;
+
+/// A hook to override the entrypoint to the main runloop used to drive the
+/// concurrency runtime and drain the main queue. This function must not return.
+/// Note: If the hook is wrapping the original function and the `compatOverride`
+///       is passed in, the `original` function pointer must be passed into the
+///       compatibility override function as the original function.
+typedef SWIFT_CC(swift) void (*swift_task_asyncMainDrainQueue_original)();
+typedef SWIFT_CC(swift) void (*swift_task_asyncMainDrainQueue_override)(
+    swift_task_asyncMainDrainQueue_original _Nullable original);
+SWIFT_EXPORT_FROM(swift_Concurrency)
+extern void *_Nullable swift_task_asyncMainDrainQueue_hook;
 
 #endif

Sources/JavaScriptKit/Runtime/index.mjs

@@ -322,6 +322,9 @@ __attribute__((__import_module__("javascript_kit"),
                __import_name__("swjs_release")))
 extern void _release(const JavaScriptObjectRef ref);
 
+__attribute__((__import_module__("javascript_kit"),
+               __import_name__("swjs_unsafe_event_loop_yield")))
+extern void _unsafe_event_loop_yield(void);
 
 #endif