Merge pull request #35215 from etcwilde/ewilde/async-main

Adding async-main support

Resolves: rdar://71828636

Author: etcwilde

include/swift/Runtime/Concurrency.h

@@ -357,6 +357,11 @@ void swift_continuation_throwingResumeWithError(/* +1 */ SwiftError *error,
 extern "C" SWIFT_CC(swift)
 void swift_continuation_logFailedCheck(const char *message);
 
+/// Drain the queue
+/// If the binary links CoreFoundation, uses CFRunLoopRun
+/// Otherwise it uses dispatchMain.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+void swift_task_asyncMainDrainQueue();
 }
 
 #endif

lib/Sema/TypeCheckAttr.cpp

@@ -1935,7 +1935,31 @@ synthesizeMainBody(AbstractFunctionDecl *fn, void *arg) {
 
   Expr *returnedExpr;
 
-  if (mainFunction->hasThrows()) {
+  if (mainFunction->hasAsync()) {
+    // Pass main into _runAsyncMain(_ asyncFunc: () async throws -> ())
+    // Resulting $main looks like:
+    // $main() { _runAsyncMain(main) }
+    auto *concurrencyModule = context.getLoadedModule(context.Id_Concurrency);
+    assert(concurrencyModule != nullptr && "Failed to find Concurrency module");
+
+    SmallVector<ValueDecl *, 1> decls;
+    concurrencyModule->lookupQualified(
+        concurrencyModule,
+        DeclNameRef(context.getIdentifier("_runAsyncMain")),
+        NL_QualifiedDefault | NL_IncludeUsableFromInline,
+        decls);
+    assert(!decls.empty() && "Failed to find _runAsyncMain");
+    FuncDecl *runner = cast<FuncDecl>(decls[0]);
+
+    auto asyncRunnerDeclRef = ConcreteDeclRef(runner, substitutionMap);
+
+    DeclRefExpr *funcExpr = new (context) DeclRefExpr(asyncRunnerDeclRef,
+                                                      DeclNameLoc(),
+                                                      /*implicit=*/true);
+    funcExpr->setType(runner->getInterfaceType());
+    auto *callExpr = CallExpr::createImplicit(context, funcExpr, memberRefExpr, {});
+    returnedExpr = callExpr;
+  } else if (mainFunction->hasThrows()) {
     auto *tryExpr = new (context) TryExpr(
         callExpr->getLoc(), callExpr, context.TheEmptyTupleType, /*implicit=*/true);
     returnedExpr = tryExpr;

stdlib/public/Concurrency/Task.cpp

@@ -22,6 +22,12 @@
 #include "TaskPrivate.h"
 #include "AsyncCall.h"
 
+#include <dispatch/dispatch.h>
+
+#if !defined(_WIN32)
+#include <dlfcn.h>
+#endif
+
 using namespace swift;
 using FutureFragment = AsyncTask::FutureFragment;
 using GroupFragment = AsyncTask::GroupFragment;
@@ -528,3 +534,19 @@ SWIFT_CC(swift)
 void swift::swift_continuation_logFailedCheck(const char *message) {
   swift_reportError(0, message);
 }
+
+void swift::swift_task_asyncMainDrainQueue() {
+#if !defined(_WIN32)
+  auto runLoop =
+      reinterpret_cast<void (*)(void)>(dlsym(RTLD_DEFAULT, "CFRunLoopRun"));
+  if (runLoop)
+    runLoop();
+  else
+    dispatch_main();
+#else
+  // TODO: I don't have a windows box to get this working right now.
+  //       We need to either pull in the CFRunLoop if it's available, or do
+  //       something that will drain the main queue. Exploding for now.
+  abort();
+#endif
+}

stdlib/public/Concurrency/Task.swift

@@ -323,6 +323,21 @@ func isTaskCancelled(_ task: Builtin.NativeObject) -> Bool
 @_silgen_name("swift_task_runAndBlockThread")
 public func runAsyncAndBlock(_ asyncFun: @escaping () async -> ())
 
+@_silgen_name("swift_task_asyncMainDrainQueue")
+public func _asyncMainDrainQueue() -> Never
+
+public func _runAsyncMain(_ asyncFun: @escaping () async throws -> ()) {
+  let _ = Task.runDetached {
+    do {
+      try await asyncFun()
+      exit(0)
+    } catch {
+      _errorInMain(error)
+    }
+  }
+  _asyncMainDrainQueue()
+}
+
 @_silgen_name("swift_task_future_wait")
 func _taskFutureWait(
   on task: Builtin.NativeObject

stdlib/public/SwiftShims/_SwiftConcurrency.h

@@ -25,6 +25,10 @@ typedef struct _SwiftContext {
   struct _SwiftContext *parentContext;
 } _SwiftContext;
 
+void exit(int);
+
+#define EXIT_SUCCESS 0
+
 #ifdef __cplusplus
 } // extern "C"
 } // namespace swift

test/Concurrency/async_main.swift

@@ -0,0 +1,38 @@
+// RUN: %target-swift-frontend -dump-ast -enable-experimental-concurrency -parse-as-library %s | %FileCheck %s --check-prefix=CHECK-AST
+// RUN: %target-build-swift -Xfrontend -enable-experimental-concurrency -Xfrontend -parse-as-library %s -o %t_binary
+// RUN: %target-run %t_binary | %FileCheck %s --check-prefix=CHECK-EXEC
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// REQUIRES: OS=macosx || OS=ios
+
+func asyncFunc() async {
+  print("Hello World!")
+}
+
+@main struct MyProgram {
+  static func main() async {
+    await asyncFunc()
+  }
+}
+
+// CHECK-EXEC: Hello World!
+
+// CHECK-AST-LABEL: "main()" interface
+// CHECK-AST:       (await_expr type='()'
+// CHECK-AST-NEXT:    (call_expr type='()'
+// CHECK-AST-NEXT:       (declref_expr type='() async -> ()'
+// CHECK-AST-SAME:        decl=async_main.(file).asyncFunc()@
+
+// CHECK-AST-LABEL: (func_decl implicit "$main()" interface
+// CHECK-AST:       (brace_stmt
+// CHECK-AST-NEXT:    (return_stmt implicit
+// CHECK-AST-NEXT:      (call_expr implicit type='()'
+// CHECK-AST-NEXT:        (declref_expr implicit
+// CHECK-AST-SAME:             type='(@escaping () async throws -> ()) -> ()'
+// CHECK-AST-SAME:             decl=_Concurrency.(file)._runAsyncMain
+// CHECK-AST-SAME:             function_ref=single
+// CHECK-AST-NEXT:        (paren_expr implicit type='(() async throws -> ())'
+// CHECK-AST-NEXT:          (function_conversion_expr implicit type='() async throws -> ()'
+// CHECK-AST-NEXT:          (dot_syntax_call_expr
+// CHECK-AST-NEXT:          (autoclosure_expr implicit type='(MyProgram.Type) -> () async -> ()'

test/Concurrency/async_main_throws_prints_error.swift

@@ -0,0 +1,22 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -Xfrontend -enable-experimental-concurrency -Xfrontend -parse-as-library %s -o %t/main
+// RUN: %target-codesign %t/main
+// RUN: %target-run %t/main  > %t/log 2>&1 || true
+// RUN: %FileCheck %s < %t/log
+
+// REQUIRES: concurrency
+// REQUIRES: executable_test
+// REQUIRES: OS=macosx || OS=ios
+
+enum Err : Error { case noGood }
+
+func asyncFunc() async throws {
+  throw Err.noGood
+}
+
+// CHECK: Fatal error: Error raised at top level: main.Err.noGood
+@main struct MyProgram {
+  static func main() async throws {
+    try await asyncFunc()
+  }
+}