Merge pull request swiftlang#36014 from DougGregor/global-actor-closures

[Concurrency] Allow async closures to be annotated with global actors.

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -4448,6 +4448,9 @@ ERROR(global_actor_non_unsafe_init,none,
 ERROR(actor_isolation_multiple_attr,none,
       "%0 %1 has multiple actor-isolation attributes ('%2' and '%3')",
       (DescriptiveDeclKind, DeclName, StringRef, StringRef))
+ERROR(global_actor_isolated_synchronous_closure,none,
+      "closure isolated to global actor %0 must be 'async'",
+      (Type))
 
 ERROR(actor_isolation_override_mismatch,none,
       "%0 %1 %2 has different actor isolation from %3 overridden declaration",

include/swift/AST/TypeCheckRequests.h

@@ -985,7 +985,8 @@ using CustomAttrNominalPair = std::pair<CustomAttr *, NominalTypeDecl *>;
 class GlobalActorAttributeRequest :
     public SimpleRequest<
         GlobalActorAttributeRequest,
-        Optional<CustomAttrNominalPair>(Decl *),
+        Optional<CustomAttrNominalPair>(
+            llvm::PointerUnion<Decl *, ClosureExpr *>),
         RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -995,7 +996,8 @@ class GlobalActorAttributeRequest :
 
   // Evaluation.
   Optional<std::pair<CustomAttr *, NominalTypeDecl *>>
-  evaluate(Evaluator &evaluator, Decl *decl) const;
+  evaluate(
+      Evaluator &evaluator, llvm::PointerUnion<Decl *, ClosureExpr *>) const;
 
 public:
   // Caching

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -102,7 +102,8 @@ SWIFT_REQUEST(TypeChecker, GlobalActorInstanceRequest,
               VarDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GlobalActorAttributeRequest,
-              Optional<CustomAttrNominalPair>(Decl *),
+              Optional<CustomAttrNominalPair>(
+                  llvm::PointerUnion<Decl *, ClosureExpr *>),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ActorIsolationRequest,
               ActorIsolationState(ValueDecl *),

lib/Sema/TypeCheckAttr.cpp

@@ -5584,9 +5584,10 @@ namespace {
 class ClosureAttributeChecker
     : public AttributeVisitor<ClosureAttributeChecker> {
   ASTContext &ctx;
+  ClosureExpr *closure;
 public:
   ClosureAttributeChecker(ClosureExpr *closure)
-    : ctx(closure->getASTContext()) { }
+    : ctx(closure->getASTContext()), closure(closure) { }
 
   void visitDeclAttribute(DeclAttribute *attr) {
     ctx.Diags.diagnose(
@@ -5599,6 +5600,29 @@ class ClosureAttributeChecker
   void visitConcurrentAttr(ConcurrentAttr *attr) {
     // Nothing else to check.
   }
+
+  void visitCustomAttr(CustomAttr *attr) {
+    // Check whether this custom attribute is the global actor attribute.
+    auto globalActorAttr = evaluateOrDefault(
+        ctx.evaluator, GlobalActorAttributeRequest{closure}, None);
+    if (globalActorAttr && globalActorAttr->first == attr)
+      return;
+
+    // Otherwise, it's an error.
+    std::string typeName;
+    if (auto typeRepr = attr->getTypeRepr()) {
+      llvm::raw_string_ostream out(typeName);
+      typeRepr->print(out);
+    } else {
+      typeName = attr->getType().getString();
+    }
+
+    ctx.Diags.diagnose(
+        attr->getLocation(), diag::unsupported_closure_attr,
+        attr->isDeclModifier(), typeName)
+      .fixItRemove(attr->getRangeWithAt());
+    attr->setInvalid();
+  }
 };
 
 }

lib/Sema/TypeCheckConcurrency.cpp

@@ -419,13 +419,26 @@ VarDecl *GlobalActorInstanceRequest::evaluate(
 
 Optional<std::pair<CustomAttr *, NominalTypeDecl *>>
 GlobalActorAttributeRequest::evaluate(
-    Evaluator &evaluator, Decl *decl) const {
-  ASTContext &ctx = decl->getASTContext();
-  auto dc = decl->getDeclContext();
+    Evaluator &evaluator,
+    llvm::PointerUnion<Decl *, ClosureExpr *> subject) const {
+  DeclContext *dc;
+  DeclAttributes *declAttrs;
+  SourceLoc loc;
+  if (auto decl = subject.dyn_cast<Decl *>()) {
+    dc = decl->getDeclContext();
+    declAttrs = &decl->getAttrs();
+    loc = decl->getLoc();
+  } else {
+    auto closure = subject.get<ClosureExpr *>();
+    dc = closure;
+    declAttrs = &closure->getAttrs();
+    loc = closure->getLoc();
+  }
+  ASTContext &ctx = dc->getASTContext();
   CustomAttr *globalActorAttr = nullptr;
   NominalTypeDecl *globalActorNominal = nullptr;
 
-  for (auto attr : decl->getAttrs().getAttributes<CustomAttr>()) {
+  for (auto attr : declAttrs->getAttributes<CustomAttr>()) {
     auto mutableAttr = const_cast<CustomAttr *>(attr);
     // Figure out which nominal declaration this custom attribute refers to.
     auto nominal = evaluateOrDefault(ctx.evaluator,
@@ -440,10 +453,10 @@ GlobalActorAttributeRequest::evaluate(
     if (!nominal->isGlobalActor())
       continue;
 
-    // Only a single global actor can be applied to a given declaration.
+    // Only a single global actor can be applied to a given entity.
     if (globalActorAttr) {
-      decl->diagnose(
-          diag::multiple_global_actors, globalActorNominal->getName(),
+      ctx.Diags.diagnose(
+          loc, diag::multiple_global_actors, globalActorNominal->getName(),
           nominal->getName());
       continue;
     }
@@ -455,8 +468,14 @@ GlobalActorAttributeRequest::evaluate(
   if (!globalActorAttr)
     return None;
 
+  // Closures can always have a global actor attached.
+  if (auto closure = subject.dyn_cast<ClosureExpr *>()) {
+    return std::make_pair(globalActorAttr, globalActorNominal);
+  }
+
   // Check that a global actor attribute makes sense on this kind of
   // declaration.
+  auto decl = subject.get<Decl *>();
   if (auto nominal = dyn_cast<NominalTypeDecl>(decl)) {
     // Nominal types are okay...
     if (auto classDecl = dyn_cast<ClassDecl>(nominal)){
@@ -1473,8 +1492,13 @@ namespace {
 
       auto dc = const_cast<DeclContext *>(constDC);
       while (!dc->isModuleScopeContext()) {
-        // Look through non-escaping closures.
         if (auto closure = dyn_cast<AbstractClosureExpr>(dc)) {
+          // If this closure has specific isolation, use it.
+          auto closureIsolation = getActorIsolationOfContext(dc);
+          if (closureIsolation != ActorIsolation::Independent)
+            return closureIsolation;
+
+          // Look through non-escaping closures.
           if (auto type = closure->getType()) {
             if (auto fnType = type->getAs<AnyFunctionType>()) {
               if (fnType->isNoEscape()) {
@@ -2039,12 +2063,50 @@ namespace {
       llvm_unreachable("unhandled actor isolation kind!");
     }
 
+    // Attempt to resolve the global actor type of a closure.
+    Type resolveGlobalActorType(ClosureExpr *closure) {
+      auto globalActorAttr = evaluateOrDefault(
+          ctx.evaluator, GlobalActorAttributeRequest{closure}, None);
+      if (!globalActorAttr)
+        return Type();
+
+      Type globalActor = evaluateOrDefault(
+          ctx.evaluator,
+          CustomAttrTypeRequest{
+            globalActorAttr->first, closure, CustomAttrTypeKind::GlobalActor},
+            Type());
+      if (!globalActor || globalActor->hasError())
+        return Type();
+
+      // Actor-isolated closures must be async.
+      bool isAsync = false;
+      if (Type closureType = closure->getType()) {
+        if (auto closureFnType = closureType->getAs<FunctionType>())
+          isAsync = closureFnType->isAsync();
+      }
+
+      if (!isAsync) {
+        ctx.Diags.diagnose(
+            closure->getLoc(), diag::global_actor_isolated_synchronous_closure,
+            globalActor);
+        return Type();
+      }
+
+      return globalActor;
+    }
+
     /// Determine the isolation of a particular closure.
     ///
     /// This function assumes that enclosing closures have already had their
     /// isolation checked.
     ClosureActorIsolation determineClosureIsolation(
         AbstractClosureExpr *closure) {
+      // If the closure specifies a global actor, use it.
+      if (auto explicitClosure = dyn_cast<ClosureExpr>(closure)) {
+        if (Type globalActorType = resolveGlobalActorType(explicitClosure))
+          return ClosureActorIsolation::forGlobalActor(globalActorType);
+      }
+
       // Escaping and concurrent closures are always actor-independent.
       if (isEscapingClosure(closure) || isConcurrentClosure(closure))
         return ClosureActorIsolation::forIndependent();

test/Concurrency/actor_isolation.swift

@@ -10,7 +10,7 @@ func acceptConcurrentClosure<T>(_: @concurrent () -> T) { }
 func acceptEscapingClosure<T>(_: @escaping () -> T) { }
 func acceptEscapingClosure<T>(_: @escaping (String) -> ()) async -> T? { nil }
 
-func acceptAsyncClosure<T>(_: () async -> T) { }
+@discardableResult func acceptAsyncClosure<T>(_: () async -> T) -> T { }
 func acceptEscapingAsyncClosure<T>(_: @escaping () async -> T) { }
 func acceptInout<T>(_: inout T) {}
 
@@ -361,6 +361,17 @@ actor Crystal {
 }
 @asyncHandler @SomeOtherGlobalActor func goo2() { await goo1() }
 
+func testGlobalActorClosures() {
+  let _: Int = acceptAsyncClosure { @SomeGlobalActor in
+    syncGlobalActorFunc()
+    syncOtherGlobalActorFunc() // expected-error{{call is 'async' but is not marked with 'await'}}
+    await syncOtherGlobalActorFunc()
+    return 17
+  }
+
+  acceptConcurrentClosure { @SomeGlobalActor in 5 } // expected-error{{closure isolated to global actor 'SomeGlobalActor' must be 'async'}}
+}
+
 extension MyActor {
   @SomeGlobalActor func onGlobalActor(otherActor: MyActor) async {
     // Access to other functions in this actor are okay.