[SILGen][Concurrency] Improve DefaultExecutorFactory lookup.

Rather than just looking at top level in the module, start by searching
the type marked as `@main`.  This means that a library that provides a
protocol or superclass that the `@main` type can conform to can specify
an executor in a reasonable manner.

Author: al45tair

lib/SILGen/SILGen.cpp

@@ -527,7 +527,36 @@ FuncDecl *SILGenModule::getExit() {
 Type SILGenModule::getConfiguredExecutorFactory() {
   auto &ctx = getASTContext();
 
-  // Look in the main module for a typealias
+  // First look in the @main struct, if any
+  NominalTypeDecl *mainType = ctx.MainModule->getMainTypeDecl();
+  if (mainType) {
+    SmallVector<ValueDecl *, 1> decls;
+    auto identifier = ctx.getIdentifier("DefaultExecutorFactory");
+    mainType->lookupQualified(mainType,
+                              DeclNameRef(identifier),
+                              SourceLoc(),
+                              NL_RemoveNonVisible | NL_RemoveOverridden
+                              | NL_OnlyTypes | NL_ProtocolMembers,
+                              decls);
+    for (auto decl : decls) {
+      TypeDecl *typeDecl = dyn_cast<TypeDecl>(decl);
+      if (typeDecl) {
+        if (auto *nominalDecl = dyn_cast<NominalTypeDecl>(typeDecl)) {
+          return nominalDecl->getDeclaredType();
+        }
+
+        if (isa<AssociatedTypeDecl>(typeDecl)) {
+          // We ignore associatedtype declarations; those with a default will
+          // turn into a `typealias` instead.
+          continue;
+        }
+
+        return typeDecl->getDeclaredInterfaceType();
+      }
+    }
+  }
+
+  // Failing that, look at the top level
   Type factory = ctx.getNamedSwiftType(ctx.MainModule, "DefaultExecutorFactory");
 
   // If we don't find it, fall back to _Concurrency.PlatformExecutorFactory

test/Concurrency/Runtime/custom_main_executor.swift

@@ -1,3 +1,8 @@
+// RUN: %target-run-simple-swift(-DTOPLEVEL_FACTORY -Xfrontend -disable-availability-checking -g %import-libdispatch -parse-as-library) | %FileCheck %s
+// RUN: %target-run-simple-swift(-DPROTOCOL_FACTORY -Xfrontend -disable-availability-checking -g %import-libdispatch -parse-as-library) | %FileCheck %s
+// RUN: %target-run-simple-swift(-DPROTOCOL_FACTORY_OVERRIDDEN -Xfrontend -disable-availability-checking -g %import-libdispatch -parse-as-library) | %FileCheck %s
+// RUN: %target-run-simple-swift(-DPROTOCOL_FACTORY_DEFAULT -Xfrontend -disable-availability-checking -g %import-libdispatch -parse-as-library) | %FileCheck %s
+// RUN: %target-run-simple-swift(-DPROTOCOL_FACTORY_DEFAULT2 -Xfrontend -disable-availability-checking -g %import-libdispatch -parse-as-library) | %FileCheck %s
 // RUN: %target-run-simple-swift(-Xfrontend -disable-availability-checking -g %import-libdispatch -parse-as-library) | %FileCheck %s
 
 // REQUIRES: concurrency
@@ -13,7 +18,9 @@
 import StdlibUnittest
 import Synchronization
 
+#if TOPLEVEL_FACTORY
 typealias DefaultExecutorFactory = SimpleExecutorFactory
+#endif
 
 struct SimpleExecutorFactory: ExecutorFactory {
   public static var mainExecutor: any MainExecutor {
@@ -26,6 +33,15 @@ struct SimpleExecutorFactory: ExecutorFactory {
   }
 }
 
+struct FatalExecutorFactory: ExecutorFactory {
+  public static var mainExecutor: any MainExecutor {
+    fatalError("mainExecutor called on FatalExecutorFactory")
+  }
+  public static var defaultExecutor: any TaskExecutor {
+    fatalError("taskExecutor called on FatalExecutorFactory")
+  }
+}
+
 @available(SwiftStdlib 6.2, *)
 final class SimpleMainExecutor: MainExecutor, @unchecked Sendable {
   public var isRunning: Bool = false
@@ -74,8 +90,30 @@ func myAsyncFunction() async {
   print("Hello World")
 }
 
+protocol AppProtocol {
+  #if PROTOCOL_FACTORY
+  associatedtype DefaultExecutorFactory
+  #endif
+  #if PROTOCOL_FACTORY_DEFAULT
+  associatedtype DefaultExecutorFactory = SimpleExecutorFactory
+  #endif
+  #if PROTOCOL_FACTORY_DEFAULT2 || PROTOCOL_FACTORY_OVERRIDDEN
+  associatedtype DefaultExecutorFactory = FatalExecutorFactory
+  #endif
+}
+
+#if PROTOCOL_FACTORY || PROTOCOL_FACTORY_DEFAULT2
+extension AppProtocol {
+  typealias DefaultExecutorFactory = SimpleExecutorFactory
+}
+#endif
+
 @available(SwiftStdlib 6.2, *)
-@main struct Main {
+@main struct Main: AppProtocol {
+  #if !TOPLEVEL_FACTORY && !PROTOCOL_FACTORY && !PROTOCOL_FACTORY_DEFAULT && !PROTOCOL_FACTORY_DEFAULT2
+  typealias DefaultExecutorFactory = SimpleExecutorFactory
+  #endif
+
   static func main() async {
     print("Hello")
     await myAsyncFunction()