Require MainActor isolation on main function

This patch forces the main function to be protected behind MainActor
isolation. If no actor isolation is specified, the main function will
implicitly have MainActor isolation.

Author: etcwilde

include/swift/AST/DiagnosticsSema.def

@@ -4289,6 +4289,10 @@ NOTE(note_add_globalactor_to_function,none,
      "add '@%0' to make %1 %2 part of global actor %3", 
      (StringRef, DescriptiveDeclKind, DeclName, Type))
 FIXIT(insert_globalactor_attr, "@%0 ", (Type))
+
+ERROR(main_function_must_be_mainActor,none,
+      "main() must be '@MainActor'", ())
+
 ERROR(not_objc_function_async,none,
       "'async' %0 cannot be represented in Objective-C", (DescriptiveDeclKind))
 NOTE(not_objc_function_type_async,none,

lib/Sema/TypeCheckConcurrency.cpp

@@ -2875,6 +2875,41 @@ static ActorIsolation getActorIsolationFromWrappedProperty(VarDecl *var) {
   return ActorIsolation::forUnspecified();
 }
 
+static Optional<ActorIsolation>
+getActorIsolationForMainFuncDecl(FuncDecl *fnDecl) {
+  // Ensure that the base type that this function is declared in has @main
+  // attribute
+  NominalTypeDecl *declContext =
+      dyn_cast<NominalTypeDecl>(fnDecl->getDeclContext());
+  if (ExtensionDecl *exDecl =
+          dyn_cast<ExtensionDecl>(fnDecl->getDeclContext())) {
+    declContext = exDecl->getExtendedNominal();
+  }
+
+  // We're not even in a nominal decl type, this can't be the main function decl
+  if (!declContext)
+    return {};
+  const bool isMainDeclContext =
+      declContext->getAttrs().hasAttribute<MainTypeAttr>();
+
+  ASTContext &ctx = fnDecl->getASTContext();
+
+  const bool isMainMain = fnDecl->isMainTypeMainMethod();
+  const bool isMain$Main =
+      fnDecl->getBaseIdentifier() == ctx.getIdentifier("$main") &&
+      !fnDecl->isInstanceMember() &&
+      fnDecl->getResultInterfaceType()->isVoid() &&
+      fnDecl->getParameters()->size() == 0;
+  const bool isMainFunction = isMainDeclContext && (isMainMain || isMain$Main);
+  const bool hasMainActor = !ctx.getMainActorType().isNull();
+
+  return isMainFunction && hasMainActor
+             ? ActorIsolation::forGlobalActor(
+                   ctx.getMainActorType()->mapTypeOutOfContext(),
+                   /*isUnsafe*/ false)
+             : Optional<ActorIsolation>();
+}
+
 /// Check rules related to global actor attributes on a class declaration.
 ///
 /// \returns true if an error occurred.
@@ -2941,9 +2976,26 @@ ActorIsolation ActorIsolationRequest::evaluate(
     return ActorIsolation::forActorInstance(actor);
   }
 
+  auto isolationFromAttr = getIsolationFromAttributes(value);
+  if (FuncDecl *fd = dyn_cast<FuncDecl>(value)) {
+    // Main.main() and Main.$main are implicitly MainActor-protected.
+    // Any other isolation is an error.
+    Optional<ActorIsolation> mainIsolation =
+        getActorIsolationForMainFuncDecl(fd);
+    if (mainIsolation) {
+      if (isolationFromAttr && isolationFromAttr->isGlobalActor()) {
+        if (!areTypesEqual(isolationFromAttr->getGlobalActor(),
+                           mainIsolation->getGlobalActor())) {
+          fd->getASTContext().Diags.diagnose(
+              fd->getLoc(), diag::main_function_must_be_mainActor);
+        }
+      }
+      return *mainIsolation;
+    }
+  }
   // If this declaration has one of the actor isolation attributes, report
   // that.
-  if (auto isolationFromAttr = getIsolationFromAttributes(value)) {
+  if (isolationFromAttr) {
     // Nonisolated declarations must involve Sendable types.
     if (*isolationFromAttr == ActorIsolation::Independent) {
       SubstitutionMap subs;

test/Concurrency/async_main.swift

@@ -10,17 +10,25 @@
 // REQUIRES: concurrency_runtime
 // UNSUPPORTED: back_deployment_runtime
 
+@MainActor
+var foo: Int = 42
+
 func asyncFunc() async {
   print("Hello World!")
 }
 
 @main struct MyProgram {
   static func main() async throws {
+    print(foo)
+    foo += 1
     await asyncFunc()
+    print(foo)
   }
 }
 
-// CHECK-EXEC: Hello World!
+// CHECK-EXEC: 42
+// CHECK-EXEC-NEXT: Hello World!
+// CHECK-EXEC-NEXT: 43
 
 // static MyProgram.main()
 // CHECK-SIL-LABEL: sil hidden @$s10async_main9MyProgramV0B0yyYaKFZ : $@convention(method) @async (@thin MyProgram.Type) -> @error Error

test/Concurrency/async_main_invalid_global_actor.swift

@@ -0,0 +1,25 @@
+// RUN: %target-typecheck-verify-swift -disable-availability-checking -parse-as-library
+
+@globalActor
+actor Kat {
+  static let shared = Kat()
+}
+
+@Kat
+var poof: Int = 1337 // expected-note{{var declared here}}
+
+@main struct Doggo {
+  @Kat
+  static func main() { // expected-error{{main() must be '@MainActor'}}
+    // expected-error@+1{{var 'poof' isolated to global actor 'Kat' can not be referenced from different global actor 'MainActor' in a synchronous context}}
+    print("Kat value: \(poof)")
+  }
+}
+
+struct Bunny {
+  // Bunnies are not @main, so they can have a "main" function that is on
+  // another actor. It's not actually the main function, so it's fine.
+  @Kat
+  static func main() {
+  }
+}

test/Concurrency/async_main_mainactor_isolation.swift

@@ -0,0 +1,13 @@
+// RUN: %target-typecheck-verify-swift -disable-availability-checking -parse-as-library
+// This should pass without any warnings or errors
+
+@MainActor
+var floofer: Int = 42
+
+@main struct Doggo { }
+
+extension Doggo {
+  static func main() {
+    print("Doggo value: \(floofer)")
+  }
+}