Add -async-main flag to favor asynchronous main

This flag biases the overload checker in favor of selecting an
asynchronous main function over a synchronous main. If no asynchronous
main function exists, a synchronous one will still be selected.
Likewise, if the flag is not passed and there are only asynchronous main
functions available, the most specific asynchronous main function will
still be selected.

Author: etcwilde

include/swift/Basic/LangOptions.h

@@ -209,6 +209,9 @@ namespace swift {
     /// Emit a remark after loading a module.
     bool EnableModuleLoadingRemarks = false;
 
+    /// Resolve main function as though it were called from an async context
+    bool EnableAsyncMainResolution = false;
+
     ///
     /// Support for alternate usage modes
     ///

include/swift/Option/Options.td

@@ -270,6 +270,10 @@ def pch_output_dir: Separate<["-"], "pch-output-dir">,
   Flags<[FrontendOption, HelpHidden, ArgumentIsPath]>,
   HelpText<"Directory to persist automatically created precompiled bridging headers">;
 
+def async_main: Flag<["-"], "async-main">,
+    Flags<[FrontendOption]>,
+    HelpText<"Resolve main function as if it were called from an asynchronous context">;
+
 // FIXME: Unhide this once it doesn't depend on an output file map.
 def incremental : Flag<["-"], "incremental">,
   Flags<[NoInteractiveOption, HelpHidden, DoesNotAffectIncrementalBuild]>,

include/swift/Sema/ConstraintSystem.h

@@ -1497,6 +1497,10 @@ enum class ConstraintSystemFlags {
   /// calling conventions, say due to Clang attributes such as
   /// `__attribute__((ns_consumed))`.
   UseClangFunctionTypes = 0x80,
+
+  /// When set, nominal typedecl contexts are asynchronous contexts.
+  /// This is set while searching for the main function
+  ConsiderNominalTypeContextsAsync = 0x100,
 };
 
 /// Options that affect the constraint system as a whole.

include/swift/Sema/IDETypeChecking.h

@@ -21,6 +21,7 @@
 
 #include "swift/AST/Identifier.h"
 #include "swift/Basic/SourceLoc.h"
+#include "swift/Sema/ConstraintSystem.h"
 #include <memory>
 #include <tuple>
 
@@ -93,8 +94,9 @@ namespace swift {
   /// Unlike other member lookup functions, \c swift::resolveValueMember()
   /// should be used when you want to look up declarations with the same name as
   /// one you already have.
-  ResolvedMemberResult resolveValueMember(DeclContext &DC, Type BaseTy,
-                                         DeclName Name);
+  ResolvedMemberResult
+  resolveValueMember(DeclContext &DC, Type BaseTy, DeclName Name,
+                     constraints::ConstraintSystemOptions Options = {});
 
   /// Given a type and an extension to the original type decl of that type,
   /// decide if the extension has been applied, i.e. if the requirements of the

lib/Driver/ToolChains.cpp

@@ -298,6 +298,7 @@ void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
                        options::OPT_verify_incremental_dependencies);
   inputArgs.AddLastArg(arguments, options::OPT_access_notes_path);
   inputArgs.AddLastArg(arguments, options::OPT_library_level);
+  inputArgs.AddLastArg(arguments, options::OPT_async_main);
 
   // Pass on any build config options
   inputArgs.AddAllArgs(arguments, options::OPT_D);

lib/Frontend/CompilerInvocation.cpp

@@ -489,6 +489,8 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
     Diags.diagnose(SourceLoc(), diag::warn_flag_deprecated,
                    "-enable-experimental-async-top-level");
 
+  Opts.EnableAsyncMainResolution = Args.hasArg(OPT_async_main);
+
   Opts.DiagnoseInvalidEphemeralnessAsError |=
       Args.hasArg(OPT_enable_invalid_ephemeralness_as_error);
 

lib/Sema/CSGen.cpp

@@ -4421,9 +4421,10 @@ getMemberDecls(InterestedMemberKind Kind) {
 }
 
 ResolvedMemberResult
-swift::resolveValueMember(DeclContext &DC, Type BaseTy, DeclName Name) {
+swift::resolveValueMember(DeclContext &DC, Type BaseTy, DeclName Name,
+                          ConstraintSystemOptions Options) {
   ResolvedMemberResult Result;
-  ConstraintSystem CS(&DC, None);
+  ConstraintSystem CS(&DC, Options);
 
   // Look up all members of BaseTy with the given Name.
   MemberLookupResult LookupResult = CS.performMemberLookup(

lib/Sema/ConstraintSystem.cpp

@@ -2875,6 +2875,11 @@ bool ConstraintSystem::isAsynchronousContext(DeclContext *dc) {
         FunctionType::ExtInfo()).isAsync();
   }
 
+  if (Options.contains(
+          ConstraintSystemFlags::ConsiderNominalTypeContextsAsync) &&
+      isa<NominalTypeDecl>(dc))
+    return true;
+
   return false;
 }
 

lib/Sema/TypeCheckAttr.cpp

@@ -2171,8 +2171,14 @@ SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
   // mainType.main() from the entry point, and that would require fully
   // type-checking the call to mainType.main().
 
+  constraints::ConstraintSystemOptions lookupOptions;
+  if (context.LangOpts.EnableAsyncMainResolution)
+    lookupOptions |=
+        constraints::ConstraintSystemFlags::ConsiderNominalTypeContextsAsync;
+
   auto resolution = resolveValueMember(
-      *declContext, nominal->getInterfaceType(), context.Id_main);
+      *declContext, nominal->getInterfaceType(), context.Id_main,
+      lookupOptions);
   FuncDecl *mainFunction =
       resolveMainFunctionDecl(declContext, resolution, context);
   if (!mainFunction) {