Merge pull request swiftlang#42142 from etcwilde/ewilde/async-main-resolution

Fixing async main resolution

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

@@ -1515,6 +1515,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/Basic/OptionSet.h"
 #include <memory>
 #include <tuple>
 
@@ -50,6 +51,8 @@ namespace swift {
   class ConstraintSystem;
   class Solution;
   class SolutionApplicationTarget;
+  enum class ConstraintSystemFlags;
+  using ConstraintSystemOptions = OptionSet<ConstraintSystemFlags>;
   }
 
   /// Typecheck binding initializer at \p bindingIndex.
@@ -93,8 +96,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

@@ -299,6 +299,7 @@ void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
   inputArgs.AddLastArg(arguments, options::OPT_access_notes_path);
   inputArgs.AddLastArg(arguments, options::OPT_library_level);
   inputArgs.AddLastArg(arguments, options::OPT_enable_regex_literals);
+  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

@@ -2894,6 +2894,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) {

test/Concurrency/async_main_resolution.swift

@@ -1,24 +1,32 @@
-// Non-apple platforms don't need to worry about the version number as much
-// because they can pull in the concurrency libraries with the swift
-// installation.
-
 // async main is nested deeper in protocols than sync, use sync
 // sync main is nested deeper in protocols than async, use async
 // async and sync are same level, use async
 
-// REQUIRES: rdar89500797
-
 // REQUIRES: concurrency
-// UNSUPPORTED: VENDOR=apple
-
-// Async is deeper in the protocol chain from `MyMain`, use sync
-// RUN: %target-swift-frontend -DASYNC_NESTED -DINHERIT_SYNC -typecheck -dump-ast -parse-as-library %s | %FileCheck %s --check-prefix=CHECK-IS-SYNC
-
-// Sync is deeper in the protocol chain from `MyMain`, use async
-// RUN: %target-swift-frontend -typecheck -dump-ast -parse-as-library %s | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
 
-// Async and sync are the same level, use async
-// RUN: %target-swift-frontend -DBOTH -DINHERIT_SYNC -typecheck -dump-ast -parse-as-library %s | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// ASYNC_NESTED: async is nested more deeply than sync main in protocol chain
+// NO_ASYNC:     no async main function
+// NO_SYNC:      no sync main function
+// BOTH:         MainProtocol has both sync and async main
+// INHERIT_SYNC: main type directly conforms to synchronous main protocol
+
+// | async flag | has async main | has sync main | both | inherits sync | nested async | Result     | Run                                                                                                                                                       |
+// |            |                |               |      |               |              | Error      | RUN: not %target-swift-frontend -disable-availability-checking -DNO_ASYNC -DNO_SYNC  -parse-as-library -typecheck -dump-ast %s 2>&1                       | %FileCheck %s --check-prefix=CHECK-IS-ERROR
+// |            | x              |               |      |               |              | Async Main | RUN: %target-swift-frontend -disable-availability-checking -DNO_SYNC -parse-as-library -typecheck -dump-ast %s                                            | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// | x          |                | x             |      |               |              | Sync Main  | RUN: %target-swift-frontend -disable-availability-checking -DNO_ASYNC -async-main -parse-as-library -typecheck -dump-ast %s                               | %FileCheck %s --check-prefix=CHECK-IS-SYNC
+// | x          | x              | x             |      |               |              | Async Main | RUN: %target-swift-frontend -disable-availability-checking -async-main -parse-as-library -typecheck -dump-ast %s                                          | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// |            | x              | x             |      |               |              | Sync Main  | RUN: %target-swift-frontend -disable-availability-checking -parse-as-library -typecheck -dump-ast %s                                                      | %FileCheck %s --check-prefix=CHECK-IS-SYNC
+// |            | x              | x             |      |               | x            | Async Main | RUN: %target-swift-frontend -disable-availability-checking -DASYNC_NESTED -parse-as-library -typecheck -dump-ast %s                                       | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// |            | x              | x             |      | x             | x            | Sync Main  | RUN: %target-swift-frontend -disable-availability-checking -DINHERIT_SYNC -DASYNC_NESTED -parse-as-library -typecheck -dump-ast %s                        | %FileCheck %s --check-prefix=CHECK-IS-SYNC
+// | x          | x              | x             |      | x             | x            | Async Main | RUN: %target-swift-frontend -disable-availability-checking -DINHERIT_SYNC -DASYNC_NESTED -async-main -parse-as-library -typecheck -dump-ast %s            | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// | x          |                | x             |      | x             | x            | Sync Main  | RUN: %target-swift-frontend -disable-availability-checking -DNO_ASYNC -DINHERIT_SYNC -DASYNC_NESTED -async-main -parse-as-library -typecheck -dump-ast %s | %FileCheck %s --check-prefix=CHECK-IS-SYNC
+// |            |                | x             | x    |               |              | Sync Main  | RUN: %target-swift-frontend -disable-availability-checking -DBOTH -DNO_ASYNC -parse-as-library -typecheck -dump-ast %s                                    | %FileCheck %s --check-prefix=CHECK-IS-SYNC
+// | x          |                | x             | x    |               |              | Async Main | RUN: %target-swift-frontend -disable-availability-checking -DBOTH -DNO_ASYNC -async-main -parse-as-library -typecheck -dump-ast %s                        | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+
+// tldr;
+// If async flag is set, will pick an asynchronous main function if one is available and related. If none exist, will fall back on synchronous main.
+// If async flag is not set, will pick a asynchronous main function if one is available and related. If none exist, will fall back on an asynchronous main
+// If neither are available; error
 
 #if ASYNC_NESTED
 protocol AsyncMainProtocol { }
@@ -55,8 +63,6 @@ extension MainProtocol {
 @main struct MyMain : AsyncMainProtocol {}
 #endif
 
-
-
 // CHECK-IS-SYNC-LABEL: "MyMain" interface type='MyMain.Type'
 // CHECK-IS-SYNC: (func_decl implicit "$main()" interface type='(MyMain.Type) -> () -> ()'
 // CHECK-IS-SYNC:       (declref_expr implicit type='(MyMain.Type) -> () -> ()'
@@ -65,6 +71,4 @@ extension MainProtocol {
 // CHECK-IS-ASYNC: (func_decl implicit "$main()" interface type='(MyMain.Type) -> () async -> ()'
 // CHECK-IS-ASYNC:       (declref_expr implicit type='(MyMain.Type) -> () async -> ()'
 
-// CHECK-IS-ERROR: error: 'MyMain' is annotated with @main and must provide a main static function of type () -> Void or () throws -> Void
-
-// CHECK-IS-ERROR-ASYNC: error: 'MyMain' is annotated with @main and must provide a main static function of type () -> Void, () throws -> Void, () async -> Void, or () async throws -> Void
+// CHECK-IS-ERROR: error: 'MyMain' is annotated with @main and must provide a main static function of type {{\(\) -> Void or \(\) throws -> Void|\(\) -> Void, \(\) throws -> Void, \(\) async -> Void, or \(\) async throws -> Void}}