Re-working async-main

include/swift/SIL/SILDeclRef.h

@@ -156,8 +156,11 @@ struct SILDeclRef {
     /// The main entry-point function. This may reference a SourceFile for a
     /// top-level main, or a decl for e.g an @main decl.
     EntryPoint,
+
+    /// The asynchronous main entry-point function.
+    AsyncEntryPoint,
   };
-  
+
   /// The AST node represented by this SILDeclRef.
   Loc loc;
   /// The Kind of this SILDeclRef.
@@ -231,6 +234,9 @@ struct SILDeclRef {
   /// Produces a SILDeclRef for a synthetic main entry-point such as @main.
   static SILDeclRef getMainDeclEntryPoint(ValueDecl *decl);
 
+  /// Produces a SILDeclRef for the synthesized async main entry-point
+  static SILDeclRef getAsyncMainDeclEntryPoint(ValueDecl *decl);
+
   /// Produces a SILDeclRef for the entry-point of a main FileUnit.
   static SILDeclRef getMainFileEntryPoint(FileUnit *file);
 

lib/IRGen/GenObjC.cpp

@@ -474,6 +474,7 @@ namespace {
       case SILDeclRef::Kind::PropertyWrapperBackingInitializer:
       case SILDeclRef::Kind::PropertyWrapperInitFromProjectedValue:
       case SILDeclRef::Kind::EntryPoint:
+      case SILDeclRef::Kind::AsyncEntryPoint:
         llvm_unreachable("Method does not have a selector");
 
       case SILDeclRef::Kind::Destroyer:

lib/SIL/IR/SILDeclRef.cpp

@@ -273,6 +273,8 @@ SILLinkage SILDeclRef::getLinkage(ForDefinition_t forDefinition) const {
   // The main entry-point is public.
   if (kind == Kind::EntryPoint)
     return SILLinkage::Public;
+  if (kind == Kind::AsyncEntryPoint)
+    return SILLinkage::Hidden;
 
   // Add External to the linkage (e.g. Public -> PublicExternal) if this is a
   // declaration not a definition.
@@ -448,6 +450,15 @@ SILDeclRef SILDeclRef::getMainDeclEntryPoint(ValueDecl *decl) {
   return result;
 }
 
+SILDeclRef SILDeclRef::getAsyncMainDeclEntryPoint(ValueDecl *decl) {
+  auto *file = cast<FileUnit>(decl->getDeclContext()->getModuleScopeContext());
+  assert(file->getMainDecl() == decl);
+  SILDeclRef result;
+  result.loc = decl;
+  result.kind = Kind::AsyncEntryPoint;
+  return result;
+}
+
 SILDeclRef SILDeclRef::getMainFileEntryPoint(FileUnit *file) {
   assert(file->hasEntryPoint() && !file->getMainDecl());
   SILDeclRef result;
@@ -539,7 +550,7 @@ IsSerialized_t SILDeclRef::isSerialized() const {
     return IsNotSerialized;
   }
 
-  if (kind == Kind::EntryPoint)
+  if (kind == Kind::EntryPoint || kind == Kind::AsyncEntryPoint)
     return IsNotSerialized;
 
   if (isIVarInitializerOrDestroyer())
@@ -925,6 +936,9 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
     return mangler.mangleInitFromProjectedValueEntity(cast<VarDecl>(getDecl()),
                                                       SKind);
 
+  case SILDeclRef::Kind::AsyncEntryPoint: {
+    return "async_Main";
+  }
   case SILDeclRef::Kind::EntryPoint: {
     return getASTContext().getEntryPointFunctionName();
   }
@@ -1257,7 +1271,8 @@ unsigned SILDeclRef::getParameterListCount() const {
     return 1;
 
   // Always uncurried even if the underlying function is curried.
-  if (kind == Kind::DefaultArgGenerator || kind == Kind::EntryPoint)
+  if (kind == Kind::DefaultArgGenerator || kind == Kind::EntryPoint ||
+      kind == Kind::AsyncEntryPoint)
     return 1;
 
   auto *vd = getDecl();

lib/SIL/IR/SILFunctionType.cpp

@@ -2526,6 +2526,9 @@ static CanSILFunctionType getNativeSILFunctionType(
     case SILDeclRef::Kind::Deallocator:
       return getSILFunctionTypeForConventions(DeallocatorConventions());
 
+    case SILDeclRef::Kind::AsyncEntryPoint:
+      return getSILFunctionTypeForConventions(
+          DefaultConventions(NormalParameterConvention::Guaranteed));
     case SILDeclRef::Kind::EntryPoint:
       llvm_unreachable("Handled by getSILFunctionTypeForAbstractCFunction");
     }
@@ -3071,6 +3074,7 @@ static ObjCSelectorFamily getObjCSelectorFamily(SILDeclRef c) {
   case SILDeclRef::Kind::PropertyWrapperBackingInitializer:
   case SILDeclRef::Kind::PropertyWrapperInitFromProjectedValue:
   case SILDeclRef::Kind::EntryPoint:
+  case SILDeclRef::Kind::AsyncEntryPoint:
     llvm_unreachable("Unexpected Kind of foreign SILDeclRef");
   }
 
@@ -3343,6 +3347,8 @@ TypeConverter::getDeclRefRepresentation(SILDeclRef c) {
     case SILDeclRef::Kind::IVarDestroyer:
       return SILFunctionTypeRepresentation::Method;
 
+    case SILDeclRef::Kind::AsyncEntryPoint:
+      return SILFunctionTypeRepresentation::Thin;
     case SILDeclRef::Kind::EntryPoint:
       return SILFunctionTypeRepresentation::CFunctionPointer;
   }
@@ -4170,6 +4176,7 @@ static AbstractFunctionDecl *getBridgedFunction(SILDeclRef declRef) {
   case SILDeclRef::Kind::IVarInitializer:
   case SILDeclRef::Kind::IVarDestroyer:
   case SILDeclRef::Kind::EntryPoint:
+  case SILDeclRef::Kind::AsyncEntryPoint:
     return nullptr;
   }
   llvm_unreachable("bad SILDeclRef kind");

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/SIL/IR/SILPrinter.cpp

@@ -316,6 +316,7 @@ void SILDeclRef::print(raw_ostream &OS) const {
   switch (kind) {
   case SILDeclRef::Kind::Func:
   case SILDeclRef::Kind::EntryPoint:
+  case SILDeclRef::Kind::AsyncEntryPoint:
     break;
   case SILDeclRef::Kind::Allocator:
     OS << "!allocator";

lib/SIL/IR/TypeLowering.cpp

@@ -2556,6 +2556,32 @@ getFunctionInterfaceTypeWithCaptures(TypeConverter &TC,
       innerExtInfo);
 }
 
+static CanAnyFunctionType getAsyncEntryPoint(ASTContext &C) {
+
+  // @main struct Main {
+  //    static func main() async throws {}
+  //    static func $main() async throws { try await main() }
+  //  }
+  //
+  // func @async_main() async -> Void {
+  //   do {
+  //      try await Main.$main()
+  //      exit(0)
+  //   } catch {
+  //      _emitErrorInMain(error)
+  //   }
+  // }
+  //
+  // This generates the type signature for @async_main
+  // TODO: 'Never' return type would be more accurate.
+
+  CanType returnType = C.getVoidType()->getCanonicalType();
+  FunctionType::ExtInfo extInfo =
+      FunctionType::ExtInfoBuilder().withAsync(true).withThrows(false).build();
+  return CanAnyFunctionType::get(/*genericSig*/ nullptr, {}, returnType,
+                                 extInfo);
+}
+
 static CanAnyFunctionType getEntryPointInterfaceType(ASTContext &C) {
   // Use standard library types if we have them; otherwise, fall back to
   // builtins.
@@ -2675,6 +2701,8 @@ CanAnyFunctionType TypeConverter::makeConstantInterfaceType(SILDeclRef c) {
   case SILDeclRef::Kind::IVarDestroyer:
     return getIVarInitDestroyerInterfaceType(cast<ClassDecl>(vd),
                                              c.isForeign, true);
+  case SILDeclRef::Kind::AsyncEntryPoint:
+    return getAsyncEntryPoint(Context);
   case SILDeclRef::Kind::EntryPoint:
     return getEntryPointInterfaceType(Context);
   }
@@ -2729,6 +2757,7 @@ TypeConverter::getConstantGenericSignature(SILDeclRef c) {
   case SILDeclRef::Kind::StoredPropertyInitializer:
     return vd->getDeclContext()->getGenericSignatureOfContext();
   case SILDeclRef::Kind::EntryPoint:
+  case SILDeclRef::Kind::AsyncEntryPoint:
     llvm_unreachable("Doesn't have generic signature");
   }
 

lib/SILGen/SILGen.cpp

@@ -338,24 +338,15 @@ SILGenModule::getConformanceToBridgedStoredNSError(SILLocation loc, Type type) {
   return SwiftModule->lookupConformance(type, proto);
 }
 
-static FuncDecl *lookupConcurrencyIntrinsic(ASTContext &C,
-                                            Optional<FuncDecl*> &cache,
-                                            StringRef name) {
+static FuncDecl *lookupIntrinsic(ModuleDecl &module,
+                                 Optional<FuncDecl *> &cache, Identifier name) {
   if (cache)
     return *cache;
-
-  auto *module = C.getLoadedModule(C.Id_Concurrency);
-  if (!module) {
-    cache = nullptr;
-    return nullptr;
-  }
-
-  SmallVector<ValueDecl *, 1> decls;
-  module->lookupQualified(module,
-                     DeclNameRef(C.getIdentifier(name)),
-                     NL_QualifiedDefault | NL_IncludeUsableFromInline,
-                     decls);
 
+  SmallVector<ValueDecl *, 1> decls;
+  module.lookupQualified(&module, DeclNameRef(name),
+                         NL_QualifiedDefault | NL_IncludeUsableFromInline,
+                         decls);
   if (decls.size() != 1) {
     cache = nullptr;
     return nullptr;
@@ -365,6 +356,18 @@ static FuncDecl *lookupConcurrencyIntrinsic(ASTContext &C,
   return func;
 }
 
+static FuncDecl *lookupConcurrencyIntrinsic(ASTContext &C,
+                                            Optional<FuncDecl *> &cache,
+                                            StringRef name) {
+  auto *module = C.getLoadedModule(C.Id_Concurrency);
+  if (!module) {
+    cache = nullptr;
+    return nullptr;
+  }
+
+  return lookupIntrinsic(*module, cache, C.getIdentifier(name));
+}
+
 FuncDecl *
 SILGenModule::getAsyncLetStart() {
   return lookupConcurrencyIntrinsic(getASTContext(),
@@ -437,6 +440,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;
@@ -1036,6 +1070,7 @@ void SILGenModule::emitFunctionDefinition(SILDeclRef constant, SILFunction *f) {
     postEmitFunction(constant, f);
     return;
   }
+  case SILDeclRef::Kind::AsyncEntryPoint:
   case SILDeclRef::Kind::EntryPoint: {
     f->setBare(IsBare);
 
@@ -1046,7 +1081,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;
   }
@@ -2029,10 +2084,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

@@ -131,6 +131,11 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   Optional<FuncDecl*> ResumeUnsafeThrowingContinuationWithError;
   Optional<FuncDecl*> CheckExpectedExecutor;
 
+  Optional<FuncDecl *> AsyncMainDrainQueue;
+  Optional<FuncDecl *> GetMainExecutor;
+  Optional<FuncDecl *> SwiftJobRun;
+  Optional<FuncDecl *> ExitFunc;
+
 public:
   SILGenModule(SILModule &M, ModuleDecl *SM);
 
@@ -517,6 +522,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);