Run first thunk of Async main synchronously

This patch updates the asynchronous main function to run the first thunk
of the function synchronously through a call to `swift_job_run`.

The runloop is killed by exiting or aborting the task that it is running
on. As such, we need to ensure that the task contains an async function
that either calls exit explicitly or aborts. The AsyncEntryPoint, that
contains this code, was added in the previous patch. This patch adds the
pieces for the actual implementation of this behaviour as well as adding
the necessary code to start the runloop.

There are now four layers of main functions before hitting the "real"
code.

@main: This is the actual main entrypoint of the program. This
constructs the task containing @async_main, grabs the main executor,
runs swift_job_run to run the first part synchronously, and finally
kicks off the runloop with a call to _asyncMainDrainQueue. This is
generated in the call to `emitAsyncMainThreadStart`.

@async_main: This thunk exists to ensure that the main function calls
`exit` at some point so that the runloop stops. It also handles emitting
an error if the user-written main function throws.

e.g:

```
func async_main() async -> () {
  do {
    try await Main.$main()
    exit(0)
  } catch {
    _errorInMain(error)
  }
}
```

Main.$main(): This still has the same behaviour as with the
synchronous case. It just calls `try await Main.main()` and exists to
simplify typechecking.

Main.main(): This is the actual user-specified main. It serves the same
purpose as in the synchronous, allowing the programmer to write code,
but it's async!

The control flow in `emitFunctionDefinition` is a little confusing (to
me anyway), so here it is spelled out:

If the main function is synchronous, the `constant.kind` will be a
`SILDeclRef::Kind::EntryPoint`, but the `decl` won't be async, so it
drops down to `emitArtificalTopLevel` anyway.

If the main function is async and we're generating `@main`, the
`constant.kind` will be `SILDeclRef::Kind::AsyncEntryPoint`, so we also
call `emitArtificalTopLevel`. `emitArtificalTopLevel` is responsible for
detecting whether the decl is async and deciding whether to emit code to
extract the argc/argv variables that get passed into the actual main
entrypoint to the program. If we're generating the `@async_main` body,
the kind will be `SILDeclRef::Kind::EntryPoint` and the `decl` will be
async, so we grab the mainEntryPoint decl and call
`emitAsyncMainThreadStart` to generate the wrapping code.

Note; there is a curious change in `SILLocation::getSourceLoc()`
where instead of simply checking `isFilenameAndLocation()`, I change it
to `getStorageKind() == FilenameAndLocationKind`. This is because the
SILLocation returned is to a FilenameAndLocationKind, but the actual
storage returns true for the call to `isNull()` inside of the
`isFilenameAndLocation()` call. This results in us incorrectly falling
through to the `getASTNode()` call below that, which asserts when asked
to get the AST node of a location.

I also did a little bit of refactoring in the SILGenModule for grabbing
intrinsics. Previously, there was only a `getConcurrencyIntrinsic`
function, which would only load FuncDecls out of the concurrency
module. The `exit` function is in the concurrency shims module, so I
refactored the load code to take a ModuleDecl to search from.

The emitBuiltinCreateAsyncTask function symbol is exposed from
SILGenBuiltin so that it is available from SILGenFunction. There is a
fair bit of work involved going from what is available at the SGF to
what is needed for actually calling the CreateAsyncTask builtin, so in
order to avoid additional maintenance, it's good to re-use that.

Author: etcwilde

lib/SIL/IR/SILLocation.cpp

@@ -42,7 +42,7 @@ SourceLoc SILLocation::getSourceLoc() const {
 
   // Don't crash if the location is a FilenameAndLocation.
   // TODO: this is a workaround until rdar://problem/25225083 is implemented.
-  if (isFilenameAndLocation())
+  if (getStorageKind() == FilenameAndLocationKind)
     return SourceLoc();
 
   return getSourceLoc(getPrimaryASTNode());

lib/SILGen/SILGen.cpp

@@ -434,6 +434,37 @@ SILGenModule::getCheckExpectedExecutor() {
                                     "_checkExpectedExecutor");
 }
 
+FuncDecl *SILGenModule::getAsyncMainDrainQueue() {
+  return lookupConcurrencyIntrinsic(getASTContext(), AsyncMainDrainQueue,
+                                    "_asyncMainDrainQueue");
+}
+
+FuncDecl *SILGenModule::getGetMainExecutor() {
+  return lookupConcurrencyIntrinsic(getASTContext(), GetMainExecutor,
+                                    "_getMainExecutor");
+}
+
+FuncDecl *SILGenModule::getSwiftJobRun() {
+  return lookupConcurrencyIntrinsic(getASTContext(), SwiftJobRun,
+                                    "_swiftJobRun");
+}
+
+FuncDecl *SILGenModule::getExit() {
+  if (ExitFunc)
+    return *ExitFunc;
+
+  ASTContext &C = getASTContext();
+  ModuleDecl *concurrencyShims =
+      C.getModuleByIdentifier(C.getIdentifier("_SwiftConcurrencyShims"));
+
+  if (!concurrencyShims) {
+    ExitFunc = nullptr;
+    return nullptr;
+  }
+
+  return lookupIntrinsic(*concurrencyShims, ExitFunc, C.getIdentifier("exit"));
+}
+
 ProtocolConformance *SILGenModule::getNSErrorConformanceToError() {
   if (NSErrorConformanceToError)
     return *NSErrorConformanceToError;
@@ -1027,7 +1058,27 @@ void SILGenModule::emitFunctionDefinition(SILDeclRef constant, SILFunction *f) {
     auto *decl = constant.getDecl();
     auto *dc = decl->getDeclContext();
     PrettyStackTraceSILFunction X("silgen emitArtificialTopLevel", f);
-    SILGenFunction(*this, *f, dc).emitArtificialTopLevel(decl);
+    // In all cases, a constant.kind == EntryPoint indicates the main entrypoint
+    // to the program, @main.
+    // In the synchronous case, the decl is not async, so emitArtificialTopLevel
+    // emits the error unwrapping and call to MainType.$main() into @main.
+    //
+    // In the async case, emitAsyncMainThreadStart is responsible for generating
+    // the contents of @main. This wraps @async_main in a task, passes that task
+    // to swift_job_run to execute the first thunk, and starts the runloop to
+    // run any additional continuations. The kind is EntryPoint, and the decl is
+    // async.
+    // When the kind is 'AsyncMain', we are generating @async_main. In this
+    // case, emitArtificialTopLevel emits the code for calling MaintType.$main,
+    // unwrapping errors, and calling exit(0) into @async_main to run the
+    // user-specified main function.
+    if (constant.kind == SILDeclRef::Kind::EntryPoint && isa<FuncDecl>(decl) &&
+        static_cast<FuncDecl *>(decl)->hasAsync()) {
+      SILDeclRef mainEntryPoint = SILDeclRef::getAsyncMainDeclEntryPoint(decl);
+      SILGenFunction(*this, *f, dc).emitAsyncMainThreadStart(mainEntryPoint);
+    } else {
+      SILGenFunction(*this, *f, dc).emitArtificialTopLevel(decl);
+    }
     postEmitFunction(constant, f);
     return;
   }
@@ -2005,10 +2056,15 @@ class SILGenModuleRAII {
 
     // If the source file contains an artificial main, emit the implicit
     // top-level code.
-    if (auto *mainDecl = sf->getMainDecl())
+    if (auto *mainDecl = sf->getMainDecl()) {
+      if (isa<FuncDecl>(mainDecl) &&
+          static_cast<FuncDecl *>(mainDecl)->hasAsync())
+        emitSILFunctionDefinition(
+            SILDeclRef::getAsyncMainDeclEntryPoint(mainDecl));
       emitSILFunctionDefinition(SILDeclRef::getMainDeclEntryPoint(mainDecl));
+    }
   }
-  
+
   void emitSILFunctionDefinition(SILDeclRef ref) {
     SGM.emitFunctionDefinition(ref, SGM.getFunction(ref, ForDefinition));
   }

lib/SILGen/SILGen.h

@@ -132,6 +132,11 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   Optional<FuncDecl*> RunAsyncHandler;
   Optional<FuncDecl*> CheckExpectedExecutor;
 
+  Optional<FuncDecl *> AsyncMainDrainQueue;
+  Optional<FuncDecl *> GetMainExecutor;
+  Optional<FuncDecl *> SwiftJobRun;
+  Optional<FuncDecl *> ExitFunc;
+
 public:
   SILGenModule(SILModule &M, ModuleDecl *SM);
 
@@ -513,6 +518,15 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Retrieve the _Concurrency._checkExpectedExecutor intrinsic.
   FuncDecl *getCheckExpectedExecutor();
 
+  /// Retrieve the _Concurrency._asyncMainDrainQueue intrinsic.
+  FuncDecl *getAsyncMainDrainQueue();
+  /// Retrieve the _Concurrency._getMainExecutor intrinsic.
+  FuncDecl *getGetMainExecutor();
+  /// Retrieve the _Concurrency._swiftJobRun intrinsic.
+  FuncDecl *getSwiftJobRun();
+  // Retrieve the _SwiftConcurrencyShims.exit intrinsic.
+  FuncDecl *getExit();
+
   SILFunction *getKeyPathProjectionCoroutine(bool isReadAccess,
                                              KeyPathTypeKind typeKind);
 

lib/SILGen/SILGenBuiltin.cpp

@@ -1402,9 +1402,10 @@ emitFunctionArgumentForAsyncTaskEntryPoint(SILGenFunction &SGF,
 }
 
 // Emit SIL for the named builtin: createAsyncTask.
-static ManagedValue emitBuiltinCreateAsyncTask(
-    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
-    ArrayRef<ManagedValue> args, SGFContext C) {
+ManagedValue emitBuiltinCreateAsyncTask(SILGenFunction &SGF, SILLocation loc,
+                                        SubstitutionMap subs,
+                                        ArrayRef<ManagedValue> args,
+                                        SGFContext C) {
   ASTContext &ctx = SGF.getASTContext();
   auto flags = args[0].forward(SGF);
 

lib/SILGen/SILGenFunction.cpp

@@ -19,6 +19,7 @@
 #include "RValue.h"
 #include "SILGenFunctionBuilder.h"
 #include "Scope.h"
+#include "swift/ABI/MetadataValues.h"
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/FileUnit.h"
@@ -609,14 +610,26 @@ void SILGenFunction::emitClosure(AbstractClosureExpr *ace) {
   emitEpilog(ace);
 }
 
+ManagedValue emitBuiltinCreateAsyncTask(SILGenFunction &SGF, SILLocation loc,
+                                        SubstitutionMap subs,
+                                        ArrayRef<ManagedValue> args,
+                                        SGFContext C);
+
 void SILGenFunction::emitArtificialTopLevel(Decl *mainDecl) {
   // Create the argc and argv arguments.
   auto entry = B.getInsertionBB();
   auto paramTypeIter = F.getConventions()
                            .getParameterSILTypes(getTypeExpansionContext())
                            .begin();
-  SILValue argc = entry->createFunctionArgument(*paramTypeIter);
-  SILValue argv = entry->createFunctionArgument(*std::next(paramTypeIter));
+
+  SILValue argc;
+  SILValue argv;
+  const bool isAsyncFunc =
+      isa<FuncDecl>(mainDecl) && static_cast<FuncDecl *>(mainDecl)->hasAsync();
+  if (!isAsyncFunc) {
+    argc = entry->createFunctionArgument(*paramTypeIter);
+    argv = entry->createFunctionArgument(*std::next(paramTypeIter));
+  }
 
   switch (mainDecl->getArtificialMainKind()) {
   case ArtificialMainKind::UIApplicationMain: {
@@ -823,13 +836,32 @@ void SILGenFunction::emitArtificialTopLevel(Decl *mainDecl) {
     auto builtinInt32Type = SILType::getBuiltinIntegerType(32, getASTContext());
 
     auto *exitBlock = createBasicBlock();
-    B.setInsertionPoint(exitBlock);
     SILValue exitCode =
         exitBlock->createPhiArgument(builtinInt32Type, OwnershipKind::None);
-    auto returnType = F.getConventions().getSingleSILResultType(B.getTypeExpansionContext());
-    if (exitCode->getType() != returnType)
-      exitCode = B.createStruct(moduleLoc, returnType, exitCode);
-    B.createReturn(moduleLoc, exitCode);
+    B.setInsertionPoint(exitBlock);
+
+    if (!mainFunc->hasAsync()) {
+      auto returnType = F.getConventions().getSingleSILResultType(
+          B.getTypeExpansionContext());
+      if (exitCode->getType() != returnType)
+        exitCode = B.createStruct(moduleLoc, returnType, exitCode);
+      B.createReturn(moduleLoc, exitCode);
+    } else {
+      FuncDecl *exitFuncDecl = SGM.getExit();
+      assert(exitFuncDecl && "Failed to find exit function declaration");
+      SILFunction *exitSILFunc = SGM.getFunction(
+          SILDeclRef(exitFuncDecl, SILDeclRef::Kind::Func, /*isForeign*/ true),
+          NotForDefinition);
+
+      SILFunctionType &funcType =
+          *exitSILFunc->getLoweredType().getAs<SILFunctionType>();
+      SILType retType = SILType::getPrimitiveObjectType(
+          funcType.getParameters().front().getInterfaceType());
+      exitCode = B.createStruct(moduleLoc, retType, exitCode);
+      SILValue exitCall = B.createFunctionRef(moduleLoc, exitSILFunc);
+      B.createApply(moduleLoc, exitCall, {}, {exitCode});
+      B.createUnreachable(moduleLoc);
+    }
 
     if (mainFunc->hasThrows()) {
       auto *successBlock = createBasicBlock();
@@ -867,6 +899,107 @@ void SILGenFunction::emitArtificialTopLevel(Decl *mainDecl) {
   }
 }
 
+void SILGenFunction::emitAsyncMainThreadStart(SILDeclRef entryPoint) {
+  auto moduleLoc = RegularLocation::getModuleLocation();
+  auto *entryBlock = B.getInsertionBB();
+  auto paramTypeIter = F.getConventions()
+                           .getParameterSILTypes(getTypeExpansionContext())
+                           .begin();
+
+  entryBlock->createFunctionArgument(*paramTypeIter);            // argc
+  entryBlock->createFunctionArgument(*std::next(paramTypeIter)); // argv
+
+  // Lookup necessary functions
+  swift::ASTContext &ctx = entryPoint.getDecl()->getASTContext();
+
+  B.setInsertionPoint(entryBlock);
+
+  /// Generates a reinterpret_cast for converting
+  /// Builtin.Job -> UnownedJob
+  /// Builtin.Executor -> UnownedSerialExecutor
+  /// These are used by _swiftJobRun, which, ABI-wise, could take
+  /// Builtin.Job or Builtin.Executor, but doesn't.
+  auto createExplodyCastForCall =
+      [this, &moduleLoc](SILValue originalValue, FuncDecl *jobRunFuncDecl,
+                         uint32_t paramIndex) -> SILValue {
+    // The type coming from the _swiftJobRun function
+    Type apiType = jobRunFuncDecl->getParameters()->get(paramIndex)->getType();
+    SILType apiSILType =
+        SILType::getPrimitiveObjectType(apiType->getCanonicalType());
+    // If the types are the same, we don't need to do anything!
+    if (apiSILType == originalValue->getType())
+      return originalValue;
+    return this->B.createUncheckedReinterpretCast(moduleLoc, originalValue,
+                                                  apiSILType);
+  };
+
+  // Call CreateAsyncTask
+  FuncDecl *builtinDecl = cast<FuncDecl>(getBuiltinValueDecl(
+      getASTContext(),
+      ctx.getIdentifier(getBuiltinName(BuiltinValueKind::CreateAsyncTask))));
+
+  auto subs = SubstitutionMap::get(builtinDecl->getGenericSignature(),
+                                   {TupleType::getEmpty(ctx)},
+                                   ArrayRef<ProtocolConformanceRef>{});
+
+  SILValue mainFunctionRef = emitGlobalFunctionRef(moduleLoc, entryPoint);
+
+  // Emit the CreateAsyncTask builtin
+  TaskCreateFlags taskCreationFlagMask(0);
+  SILValue taskFlags =
+      emitWrapIntegerLiteral(moduleLoc, getLoweredType(ctx.getIntType()),
+                             taskCreationFlagMask.getOpaqueValue());
+
+  SILValue task =
+      emitBuiltinCreateAsyncTask(*this, moduleLoc, subs,
+                                 {ManagedValue::forUnmanaged(taskFlags),
+                                  ManagedValue::forUnmanaged(mainFunctionRef)},
+                                 {})
+          .forward(*this);
+  DestructureTupleInst *structure = B.createDestructureTuple(moduleLoc, task);
+  task = structure->getResult(0);
+
+  // Get swiftJobRun
+  FuncDecl *swiftJobRunFuncDecl = SGM.getSwiftJobRun();
+  SILFunction *swiftJobRunSILFunc =
+      SGM.getFunction(SILDeclRef(swiftJobRunFuncDecl, SILDeclRef::Kind::Func),
+                      NotForDefinition);
+  SILValue swiftJobRunFunc =
+      B.createFunctionRefFor(moduleLoc, swiftJobRunSILFunc);
+
+  // Convert task to job
+  SILType JobType = SILType::getPrimitiveObjectType(
+      getBuiltinType(ctx, "Job")->getCanonicalType());
+  SILValue jobResult = B.createBuiltin(
+      moduleLoc,
+      ctx.getIdentifier(getBuiltinName(BuiltinValueKind::ConvertTaskToJob)),
+      JobType, {}, {task});
+  jobResult = createExplodyCastForCall(jobResult, swiftJobRunFuncDecl, 0);
+
+  // Get main executor
+  FuncDecl *getMainExecutorFuncDecl = SGM.getGetMainExecutor();
+  SILFunction *getMainExeutorSILFunc = SGM.getFunction(
+      SILDeclRef(getMainExecutorFuncDecl, SILDeclRef::Kind::Func),
+      NotForDefinition);
+  SILValue getMainExeutorFunc =
+      B.createFunctionRefFor(moduleLoc, getMainExeutorSILFunc);
+  SILValue mainExecutor = B.createApply(moduleLoc, getMainExeutorFunc, {}, {});
+  mainExecutor = createExplodyCastForCall(mainExecutor, swiftJobRunFuncDecl, 1);
+
+  // Run first part synchronously
+  B.createApply(moduleLoc, swiftJobRunFunc, {}, {jobResult, mainExecutor});
+
+  // Start Main loop!
+  FuncDecl *drainQueueFuncDecl = SGM.getAsyncMainDrainQueue();
+  SILFunction *drainQueueSILFunc = SGM.getFunction(
+      SILDeclRef(drainQueueFuncDecl, SILDeclRef::Kind::Func), NotForDefinition);
+  SILValue drainQueueFunc =
+      B.createFunctionRefFor(moduleLoc, drainQueueSILFunc);
+  B.createApply(moduleLoc, drainQueueFunc, {}, {});
+  B.createUnreachable(moduleLoc);
+  return;
+}
+
 void SILGenFunction::emitGeneratorFunction(SILDeclRef function, Expr *value,
                                            bool EmitProfilerIncrement) {
   auto *dc = function.getDecl()->getInnermostDeclContext();

lib/Sema/TypeCheckAttr.cpp

@@ -1888,39 +1888,29 @@ synthesizeMainBody(AbstractFunctionDecl *fn, void *arg) {
   Expr *returnedExpr;
 
   if (mainFunction->hasAsync()) {
-    // Pass main into _runAsyncMain(_ asyncFunc: () async throws -> ())
-    // Resulting $main looks like:
-    // $main() { _runAsyncMain(main) }
+    // Ensure that the concurrency module is loaded
     auto *concurrencyModule = context.getLoadedModule(context.Id_Concurrency);
     if (!concurrencyModule) {
       context.Diags.diagnose(mainFunction->getAsyncLoc(),
                              diag::async_main_no_concurrency);
       auto result = new (context) ErrorExpr(mainFunction->getSourceRange());
-      SmallVector<ASTNode, 1> stmts;
-      stmts.push_back(result);
-      auto body = BraceStmt::create(context, SourceLoc(), stmts,
-                                    SourceLoc(), /*Implicit*/true);
-
+      ASTNode stmts[] = {result};
+      auto body = BraceStmt::create(context, SourceLoc(), stmts, SourceLoc(),
+                                    /*Implicit*/ true);
       return std::make_pair(body, /*typechecked*/true);
     }
 
-    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;
+    // $main() async { await main() }
+    Expr *awaitExpr =
+        new (context) AwaitExpr(callExpr->getLoc(), callExpr,
+                                context.TheEmptyTupleType, /*implicit*/ true);
+    if (mainFunction->hasThrows()) {
+      // $main() async throws { try await main() }
+      awaitExpr =
+          new (context) TryExpr(callExpr->getLoc(), awaitExpr,
+                                context.TheEmptyTupleType, /*implicit*/ true);
+    }
+    returnedExpr = awaitExpr;
   } else if (mainFunction->hasThrows()) {
     auto *tryExpr = new (context) TryExpr(
         callExpr->getLoc(), callExpr, context.TheEmptyTupleType, /*implicit=*/true);
@@ -2037,7 +2027,7 @@ SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
       DeclName(context, DeclBaseName(context.Id_MainEntryPoint),
                ParameterList::createEmpty(context)),
       /*NameLoc=*/SourceLoc(),
-      /*Async=*/false,
+      /*Async=*/mainFunction->hasAsync(),
       /*Throws=*/mainFunction->hasThrows(),
       /*GenericParams=*/nullptr, ParameterList::createEmpty(context),
       /*FnRetType=*/TupleType::getEmpty(context), declContext);
@@ -2498,7 +2488,7 @@ static FuncDecl *findSimilarAccessor(DeclNameRef replacedVarName,
   }
 
   assert(!isa<FuncDecl>(results[0]));
-  
+
   auto *origStorage = cast<AbstractStorageDecl>(results[0]);
   if (forDynamicReplacement && !origStorage->isDynamic()) {
     Diags.diagnose(attr->getLocation(),