Use Disjunction Constraint to find main function

include/swift/Basic/LangOptions.h

@@ -223,9 +223,6 @@ 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,10 +270,6 @@ 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

@@ -1535,7 +1535,7 @@ enum class ConstraintSystemFlags {
   /// Note that this flag is automatically applied to all constraint systems,
   /// when \c DebugConstraintSolver is set in \c TypeCheckerOptions. It can be
   /// automatically enabled for select constraint solving attempts by setting
-  /// \c DebugConstraintSolverAttempt. Finally, it can also be automatically 
+  /// \c DebugConstraintSolverAttempt. Finally, it can also be automatically
   /// enabled for a pre-configured set of expressions on line numbers by setting
   /// \c DebugConstraintSolverOnLines.
   DebugConstraints = 0x10,
@@ -1560,9 +1560,8 @@ enum class ConstraintSystemFlags {
   /// `__attribute__((ns_consumed))`.
   UseClangFunctionTypes = 0x80,
 
-  /// When set, nominal typedecl contexts are asynchronous contexts.
-  /// This is set while searching for the main function
-  ConsiderNominalTypeContextsAsync = 0x100,
+  /// When set, ignore async/sync mismatches
+  IgnoreAsyncSyncMismatch = 0x100,
 };
 
 /// Options that affect the constraint system as a whole.

include/swift/Sema/IDETypeChecking.h

@@ -21,7 +21,6 @@
 
 #include "swift/AST/Identifier.h"
 #include "swift/Basic/SourceLoc.h"
-#include "swift/Basic/OptionSet.h"
 #include <memory>
 #include <tuple>
 
@@ -51,8 +50,6 @@ namespace swift {
   class ConstraintSystem;
   class Solution;
   class SolutionApplicationTarget;
-  enum class ConstraintSystemFlags;
-  using ConstraintSystemOptions = OptionSet<ConstraintSystemFlags>;
   }
 
   /// Typecheck binding initializer at \p bindingIndex.
@@ -96,9 +93,8 @@ 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,
-                     constraints::ConstraintSystemOptions Options = {});
+  ResolvedMemberResult resolveValueMember(DeclContext &DC, Type BaseTy,
+                                          DeclName Name);
 
   /// 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

@@ -300,7 +300,6 @@ 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_bare_slash_regex);
-  inputArgs.AddLastArg(arguments, options::OPT_async_main);
 
   // Pass on any build config options
   inputArgs.AddAllArgs(arguments, options::OPT_D);

lib/Frontend/CompilerInvocation.cpp

@@ -497,8 +497,6 @@ 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

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

lib/Sema/ConstraintSystem.cpp

@@ -2900,11 +2900,6 @@ bool ConstraintSystem::isAsynchronousContext(DeclContext *dc) {
         FunctionType::ExtInfo()).isAsync();
   }
 
-  if (Options.contains(
-          ConstraintSystemFlags::ConsiderNominalTypeContextsAsync) &&
-      isa<NominalTypeDecl>(dc))
-    return true;
-
   return false;
 }
 
@@ -3309,7 +3304,8 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     // If we're choosing an asynchronous declaration within a synchronous
     // context, or vice-versa, increase the async/async mismatch score.
     if (auto func = dyn_cast<AbstractFunctionDecl>(decl)) {
-      if (!func->hasPolymorphicEffect(EffectKind::Async) &&
+      if (!Options.contains(ConstraintSystemFlags::IgnoreAsyncSyncMismatch) &&
+          !func->hasPolymorphicEffect(EffectKind::Async) &&
           func->isAsyncContext() != isAsynchronousContext(useDC)) {
         increaseScore(
             func->isAsyncContext() ? SK_AsyncInSyncMismatch : SK_SyncInAsync);

lib/Sema/TypeCheckAttr.cpp

@@ -2096,28 +2096,6 @@ synthesizeMainBody(AbstractFunctionDecl *fn, void *arg) {
   return std::make_pair(body, /*typechecked=*/false);
 }
 
-static FuncDecl *resolveMainFunctionDecl(DeclContext *declContext,
-                                         ResolvedMemberResult &resolution,
-                                         ASTContext &ctx) {
-  // Choose the best overload if it's a main function
-  if (resolution.hasBestOverload()) {
-    ValueDecl *best = resolution.getBestOverload();
-    if (FuncDecl *func = dyn_cast<FuncDecl>(best)) {
-      if (func->isMainTypeMainMethod()) {
-        return func;
-      }
-    }
-  }
-  // Look for the most highly-ranked main-function candidate
-  for (ValueDecl *candidate : resolution.getMemberDecls(Viable)) {
-    if (FuncDecl *func = dyn_cast<FuncDecl>(candidate)) {
-      if (func->isMainTypeMainMethod())
-        return func;
-    }
-  }
-  return nullptr;
-}
-
 FuncDecl *
 SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
                                         Decl *D) const {
@@ -2170,17 +2148,91 @@ SynthesizeMainFunctionRequest::evaluate(Evaluator &evaluator,
   // usual type-checking.  The alternative would be to directly call
   // mainType.main() from the entry point, and that would require fully
   // type-checking the call to mainType.main().
+  using namespace constraints;
+  ConstraintSystem CS(declContext,
+                      ConstraintSystemFlags::IgnoreAsyncSyncMismatch);
+  ConstraintLocator *locator =
+      CS.getConstraintLocator({}, ConstraintLocator::Member);
+  // Allowed main function types
+  // `() -> Void`
+  // `() async -> Void`
+  // `() throws -> Void`
+  // `() async throws -> Void`
+  // `@MainActor () -> Void`
+  // `@MainActor () async -> Void`
+  // `@MainActor () throws -> Void`
+  // `@MainActor () async throws -> Void`
+  {
+    llvm::SmallVector<Type, 8> mainTypes = {
+
+        FunctionType::get(/*params*/ {}, context.TheEmptyTupleType,
+                          ASTExtInfo()),
+        FunctionType::get(
+            /*params*/ {}, context.TheEmptyTupleType,
+            ASTExtInfoBuilder().withAsync().build()),
+
+        FunctionType::get(/*params*/ {}, context.TheEmptyTupleType,
+                          ASTExtInfoBuilder().withThrows().build()),
+
+        FunctionType::get(
+            /*params*/ {}, context.TheEmptyTupleType,
+            ASTExtInfoBuilder().withAsync().withThrows().build())};
+
+    Type mainActor = context.getMainActorType();
+    if (mainActor) {
+      mainTypes.push_back(FunctionType::get(
+          /*params*/ {}, context.TheEmptyTupleType,
+          ASTExtInfoBuilder().withGlobalActor(mainActor).build()));
+      mainTypes.push_back(FunctionType::get(
+          /*params*/ {}, context.TheEmptyTupleType,
+          ASTExtInfoBuilder().withAsync().withGlobalActor(mainActor).build()));
+      mainTypes.push_back(FunctionType::get(
+          /*params*/ {}, context.TheEmptyTupleType,
+          ASTExtInfoBuilder().withThrows().withGlobalActor(mainActor).build()));
+      mainTypes.push_back(FunctionType::get(/*params*/ {},
+                                            context.TheEmptyTupleType,
+                                            ASTExtInfoBuilder()
+                                                .withAsync()
+                                                .withThrows()
+                                                .withGlobalActor(mainActor)
+                                                .build()));
+    }
+    TypeVariableType *mainType =
+        CS.createTypeVariable(locator, /*options=*/0);
+    llvm::SmallVector<Constraint *, 4> typeEqualityConstraints;
+    typeEqualityConstraints.reserve(mainTypes.size());
+    for (const Type &candidateMainType : mainTypes) {
+      typeEqualityConstraints.push_back(
+          Constraint::create(CS, ConstraintKind::Equal, Type(mainType),
+                             candidateMainType, locator));
+    }
+
+    CS.addDisjunctionConstraint(typeEqualityConstraints, locator);
+    CS.addValueMemberConstraint(
+        nominal->getInterfaceType(), DeclNameRef(context.Id_main),
+        Type(mainType), declContext, FunctionRefKind::SingleApply, {}, locator);
+  }
+
+  FuncDecl *mainFunction = nullptr;
+  llvm::SmallVector<Solution, 4> candidates;
+
+  if (!CS.solve(candidates, FreeTypeVariableBinding::Disallow)) {
+    // We can't use CS.diagnoseAmbiguity directly since the locator is empty
+    // Sticking the main type decl `D` in results in an assert due to a
+    // unsimplifiable locator anchor since it appears to be looking for an
+    // expression, which we don't have.
+    // (locator could not be simplified to anchor)
+    // TODO: emit notes for each of the ambiguous candidates
+    if (candidates.size() != 1) {
+      context.Diags.diagnose(nominal->getLoc(), diag::ambiguous_decl_ref,
+                             DeclNameRef(context.Id_main));
+      attr->setInvalid();
+      return nullptr;
+    }
+    mainFunction = dyn_cast<FuncDecl>(
+        candidates[0].overloadChoices[locator].choice.getDecl());
+  }
 
-  constraints::ConstraintSystemOptions lookupOptions;
-  if (context.LangOpts.EnableAsyncMainResolution)
-    lookupOptions |=
-        constraints::ConstraintSystemFlags::ConsiderNominalTypeContextsAsync;
-
-  auto resolution = resolveValueMember(
-      *declContext, nominal->getInterfaceType(), context.Id_main,
-      lookupOptions);
-  FuncDecl *mainFunction =
-      resolveMainFunctionDecl(declContext, resolution, context);
   if (!mainFunction) {
     const bool hasAsyncSupport =
         AvailabilityContext::forDeploymentTarget(context).isContainedIn(

lib/Sema/TypeCheckConcurrency.cpp

@@ -3380,7 +3380,8 @@ getActorIsolationForMainFuncDecl(FuncDecl *fnDecl) {
   if (!declContext)
     return {};
   const bool isMainDeclContext =
-      declContext->getAttrs().hasAttribute<MainTypeAttr>();
+      declContext->getAttrs().hasAttribute<MainTypeAttr>(
+          /*allow invalid*/ true);
 
   ASTContext &ctx = fnDecl->getASTContext();
 

test/Concurrency/async_main_resolution.swift

@@ -1,6 +1,11 @@
+// This test aims to show that no preference is given to either the async or
+// sync main function. The most specific, valid, main function will be
+// selected if one exists. If two main functions could exist, the usage is
+// ambiguous.
+
 // 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
+// async and sync are same level, error
 
 // REQUIRES: concurrency
 
@@ -10,23 +15,17 @@
 // 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
+// | has async main | has sync main | both | inherits sync | nested async | Result |                         | Run                                                                                                                                          |
+// |                |               |      |               |              | Error  | No main                 | RUN: not %target-swift-frontend -disable-availability-checking -DNO_SYNC -DNO_ASYNC -parse-as-library -typecheck -dump-ast %s 2>&1           | %FileCheck %s --check-prefix=CHECK-IS-ERROR1
+// | x              | x             | x    | x             |              | Error  | Ambiguous main in MainP | RUN: not %target-swift-frontend -disable-availability-checking -DBOTH -DINHERIT_SYNC -parse-as-library -typecheck -dump-ast %s 2>&1          | %FileCheck %s --check-prefix=CHECK-IS-ERROR2
+// |                | x             | x    | x             |              | Error  | Ambiguous main in MainP | RUN: not %target-swift-frontend -disable-availability-checking -DBOTH -DINHERIT_SYNC -parse-as-library -typecheck -dump-ast %s 2>&1          | %FileCheck %s --check-prefix=CHECK-IS-ERROR2
+// | x              | x             | x    |               |              | Async  | Directly selected       | RUN: %target-swift-frontend -disable-availability-checking -DBOTH -parse-as-library -typecheck -dump-ast %s                                  | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// | x              | x             |      |               |              | Async  | Directly selected       | RUN: %target-swift-frontend -disable-availability-checking -parse-as-library -typecheck -dump-ast %s                                         | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// |                | x             |      |               |              | Sync   | Indirectly selected     | RUN: %target-swift-frontend -disable-availability-checking -DNO_ASYNC -parse-as-library -typecheck -dump-ast %s                              | %FileCheck %s --check-prefix=CHECK-IS-SYNC
+// | x              | x             |      | x             | x            | Sync   | Directly selected       | 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            | Async  | Indirectly selected     | RUN: %target-swift-frontend -disable-availability-checking -DNO_SYNC -DINHERIT_SYNC -DASYNC_NESTED -parse-as-library -typecheck -dump-ast %s | %FileCheck %s --check-prefix=CHECK-IS-ASYNC
+// | x              |               |      | x             |              | Error  | Unrelated async main    | RUN: not %target-swift-frontend -disable-availability-checking -DNO_SYNC -DINHERIT_SYNC -parse-as-library -typecheck -dump-ast %s 2>&1       | %FileCheck %s --check-prefix=CHECK-IS-ERROR1
+// |                | x             |      |               | x            | Error  | Unrelated sync main     | RUN: not %target-swift-frontend -disable-availability-checking -DNO_ASYNC -DASYNC_NESTED -parse-as-library -typecheck -dump-ast %s  2>&1     | %FileCheck %s --check-prefix=CHECK-IS-ERROR1
 
 #if ASYNC_NESTED
 protocol AsyncMainProtocol { }
@@ -71,4 +70,5 @@ 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|\(\) -> Void, \(\) throws -> Void, \(\) async -> Void, or \(\) async throws -> Void}}
+// CHECK-IS-ERROR1: 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}}
+// CHECK-IS-ERROR2: error: ambiguous use of 'main'