[Concurrency] Add @_inheritActorContext hidden parameter attribute.

This new attribute can be used on parameters of `@Sendable async` type
to indicate that the closures arguments passed to such parameters
should inherit the actor context where they are formed, which is not
the normal behavior for `@Sendable` closures.

Another part of rdar://76927008.

Author: DougGregor

include/swift/AST/Attr.def

@@ -652,6 +652,11 @@ SIMPLE_DECL_ATTR(_implicitSelfCapture, ImplicitSelfCapture,
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIBreakingToRemove,
   115)
 
+SIMPLE_DECL_ATTR(_inheritActorContext, InheritActorContext,
+  OnParam | UserInaccessible |
+  ABIStableToAdd | ABIStableToRemove | APIBreakingToAdd | APIBreakingToRemove,
+  116)
+
 #undef TYPE_ATTR
 #undef DECL_ATTR_ALIAS
 #undef CONTEXTUAL_DECL_ATTR_ALIAS

include/swift/AST/Expr.h

@@ -292,14 +292,18 @@ class alignas(8) Expr {
     Kind : 2
   );
 
-  SWIFT_INLINE_BITFIELD(ClosureExpr, AbstractClosureExpr, 1+1,
+  SWIFT_INLINE_BITFIELD(ClosureExpr, AbstractClosureExpr, 1+1+1,
     /// True if closure parameters were synthesized from anonymous closure
     /// variables.
     HasAnonymousClosureVars : 1,
 
     /// True if "self" can be captured implicitly without requiring "self."
     /// on each member reference.
-    ImplicitSelfCapture : 1
+    ImplicitSelfCapture : 1,
+
+    /// True if this @Sendable async closure parameter should implicitly
+    /// inherit the actor context from where it was formed.
+    InheritActorContext : 1
   );
 
   SWIFT_INLINE_BITFIELD_FULL(BindOptionalExpr, Expr, 16,
@@ -3876,6 +3880,7 @@ class ClosureExpr : public AbstractClosureExpr {
     setParameterList(params);
     Bits.ClosureExpr.HasAnonymousClosureVars = false;
     Bits.ClosureExpr.ImplicitSelfCapture = false;
+    Bits.ClosureExpr.InheritActorContext = false;
   }
 
   SourceRange getSourceRange() const;
@@ -3914,6 +3919,16 @@ class ClosureExpr : public AbstractClosureExpr {
     Bits.ClosureExpr.ImplicitSelfCapture = value;
   }
 
+  /// Whether this closure should implicitly inherit the actor context from
+  /// where it was formed. This only affects @Sendable async closures.
+  bool inheritsActorContext() const {
+    return Bits.ClosureExpr.InheritActorContext;
+  }
+
+  void setInheritsActorContext(bool value = true) {
+    Bits.ClosureExpr.InheritActorContext = value;
+  }
+
   /// Determine whether this closure expression has an
   /// explicitly-specified result type.
   bool hasExplicitResultType() const { return ArrowLoc.isValid(); }

include/swift/AST/Types.h

@@ -3382,6 +3382,7 @@ struct ParameterListInfo {
   SmallBitVector unsafeSendable;
   SmallBitVector unsafeMainActor;
   SmallBitVector implicitSelfCapture;
+  SmallBitVector inheritActorContext;
 
 public:
   ParameterListInfo() { }
@@ -3414,6 +3415,10 @@ struct ParameterListInfo {
   /// 'self' to be implicit, without requiring "self.".
   bool isImplicitSelfCapture(unsigned paramIdx) const;
 
+  /// Whether the given parameter is a closure that should inherit the
+  /// actor context from the context in which it was created.
+  bool inheritsActorContext(unsigned paramIdx) const;
+
   /// Whether there is any contextual information set on this parameter list.
   bool anyContextualInfo() const;
 

lib/AST/ASTDumper.cpp

@@ -2548,6 +2548,8 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
       PrintWithColorRAII(OS, ClosureModifierColor) << " single-expression";
     if (E->allowsImplicitSelfCapture())
       PrintWithColorRAII(OS, ClosureModifierColor) << " implicit-self";
+    if (E->inheritsActorContext())
+      PrintWithColorRAII(OS, ClosureModifierColor) << " inherits-actor-context";
 
     if (E->getParameters()) {
       OS << '\n';

lib/AST/Type.cpp

@@ -962,6 +962,7 @@ ParameterListInfo::ParameterListInfo(
   unsafeSendable.resize(params.size());
   unsafeMainActor.resize(params.size());
   implicitSelfCapture.resize(params.size());
+  inheritActorContext.resize(params.size());
 
   // No parameter owner means no parameter list means no default arguments
   // - hand back the zeroed bitvector.
@@ -1025,6 +1026,10 @@ ParameterListInfo::ParameterListInfo(
     if (param->getAttrs().hasAttribute<ImplicitSelfCaptureAttr>()) {
       implicitSelfCapture.set(i);
     }
+
+    if (param->getAttrs().hasAttribute<InheritActorContextAttr>()) {
+      inheritActorContext.set(i);
+    }
   }
 }
 
@@ -1061,9 +1066,15 @@ bool ParameterListInfo::isImplicitSelfCapture(unsigned paramIdx) const {
       : false;
 }
 
+bool ParameterListInfo::inheritsActorContext(unsigned paramIdx) const {
+  return paramIdx < inheritActorContext.size()
+      ? inheritActorContext[paramIdx]
+      : false;
+}
+
 bool ParameterListInfo::anyContextualInfo() const {
   return unsafeSendable.any() || unsafeMainActor.any() ||
-      implicitSelfCapture.any();
+      implicitSelfCapture.any() || inheritActorContext.any();
 }
 
 /// Turn a param list into a symbolic and printable representation that does not

lib/Sema/CSApply.cpp

@@ -5710,17 +5710,19 @@ static bool hasCurriedSelf(ConstraintSystem &cs, ConcreteDeclRef callee,
 
 /// Apply the contextually Sendable flag to the given expression,
 static void applyContextualClosureFlags(
-      Expr *expr, bool sendable, bool forMainActor, bool implicitSelfCapture) {
+      Expr *expr, bool sendable, bool forMainActor, bool implicitSelfCapture,
+      bool inheritActorContext) {
   if (auto closure = dyn_cast<ClosureExpr>(expr)) {
     closure->setUnsafeConcurrent(sendable, forMainActor);
     closure->setAllowsImplicitSelfCapture(implicitSelfCapture);
+    closure->setInheritsActorContext(inheritActorContext);
     return;
   }
 
   if (auto captureList = dyn_cast<CaptureListExpr>(expr)) {
     applyContextualClosureFlags(
         captureList->getClosureBody(), sendable, forMainActor,
-        implicitSelfCapture);
+        implicitSelfCapture, inheritActorContext);
   }
 }
 
@@ -5958,9 +5960,10 @@ Expr *ExprRewriter::coerceCallArguments(
     bool isMainActor = paramInfo.isUnsafeMainActor(paramIdx) ||
         (isUnsafeSendable && apply && isMainDispatchQueue(apply->getFn()));
     bool isImplicitSelfCapture = paramInfo.isImplicitSelfCapture(paramIdx);
+    bool inheritsActorContext = paramInfo.inheritsActorContext(paramIdx);
     applyContextualClosureFlags(
         arg, isUnsafeSendable && contextUsesConcurrencyFeatures(dc),
-        isMainActor, isImplicitSelfCapture);
+        isMainActor, isImplicitSelfCapture, inheritsActorContext);
 
     // If the types exactly match, this is easy.
     auto paramType = param.getOldType();
@@ -6976,8 +6979,8 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
       }
     }
 
-    // If we have a ClosureExpr, then we can safely propagate the 'concurrent'
-    // bit to the closure without invalidating prior analysis.
+    // If we have a ClosureExpr, then we can safely propagate @Sendable
+    // to the closure without invalidating prior analysis.
     auto fromEI = fromFunc->getExtInfo();
     if (toEI.isSendable() && !fromEI.isSendable()) {
       auto newFromFuncType = fromFunc->withExtInfo(fromEI.withConcurrent());

lib/Sema/TypeCheckAttr.cpp

@@ -138,6 +138,7 @@ class AttributeChecker : public AttributeVisitor<AttributeChecker> {
   IGNORED_ATTR(UnsafeSendable)
   IGNORED_ATTR(UnsafeMainActor)
   IGNORED_ATTR(ImplicitSelfCapture)
+  IGNORED_ATTR(InheritActorContext)
 #undef IGNORED_ATTR
 
   void visitAlignmentAttr(AlignmentAttr *attr) {

lib/Sema/TypeCheckConcurrency.cpp

@@ -773,19 +773,26 @@ static bool isAsyncCall(const ApplyExpr *call) {
 /// features.
 static bool shouldDiagnoseExistingDataRaces(const DeclContext *dc);
 
-/// Determine whether this closure is escaping.
-static bool isSendableClosure(const AbstractClosureExpr *closure) {
+/// Determine whether this closure should be treated as Sendable.
+///
+/// \param forActorIsolation Whether this check is for the purposes of
+/// determining whether the closure must be non-isolated.
+static bool isSendableClosure(
+    const AbstractClosureExpr *closure, bool forActorIsolation) {
+  if (auto explicitClosure = dyn_cast<ClosureExpr>(closure)) {
+    if (forActorIsolation && explicitClosure->inheritsActorContext())
+      return false;
+
+    if (explicitClosure->isUnsafeSendable())
+      return true;
+  }
+
   if (auto type = closure->getType()) {
     if (auto fnType = type->getAs<AnyFunctionType>())
       if (fnType->isSendable())
         return true;
   }
 
-  if (auto explicitClosure = dyn_cast<ClosureExpr>(closure)) {
-    if (explicitClosure->isUnsafeSendable())
-      return true;
-  }
-
   return false;
 }
 
@@ -2113,7 +2120,7 @@ namespace {
       }
 
       if (auto closure = dyn_cast<AbstractClosureExpr>(dc)) {
-        if (isSendableClosure(closure)) {
+        if (isSendableClosure(closure, /*forActorIsolation=*/true)) {
           return diag::actor_isolated_from_concurrent_closure;
         }
 
@@ -2310,8 +2317,9 @@ namespace {
         }
       }
 
-      // Sendable closures are always actor-independent.
-      if (isSendableClosure(closure))
+      // Sendable closures are actor-independent unless the closure has
+      // specifically opted into inheriting actor isolation.
+      if (isSendableClosure(closure, /*forActorIsolation=*/true))
         return ClosureActorIsolation::forIndependent();
 
       // A non-escaping closure gets its isolation from its context.
@@ -2362,7 +2370,7 @@ bool ActorIsolationChecker::mayExecuteConcurrentlyWith(
   while (useContext != defContext) {
     // If we find a concurrent closure... it can be run concurrently.
     if (auto closure = dyn_cast<AbstractClosureExpr>(useContext)) {
-      if (isSendableClosure(closure))
+      if (isSendableClosure(closure, /*forActorIsolation=*/false))
         return true;
     }
 

lib/Sema/TypeCheckDeclOverride.cpp

@@ -1547,6 +1547,7 @@ namespace  {
     UNINTERESTING_ATTR(UnsafeSendable)
     UNINTERESTING_ATTR(UnsafeMainActor)
     UNINTERESTING_ATTR(ImplicitSelfCapture)
+    UNINTERESTING_ATTR(InheritActorContext)
 #undef UNINTERESTING_ATTR
 
     void visitAvailableAttr(AvailableAttr *attr) {

test/Concurrency/actor_isolation.swift

@@ -70,7 +70,7 @@ actor MyActor: MySuperActor {
   class func synchronousClass() { }
   static func synchronousStatic() { }
 
-  func synchronous() -> String { text.first ?? "nothing" } // expected-note 19{{calls to instance method 'synchronous()' from outside of its actor context are implicitly asynchronous}}
+  func synchronous() -> String { text.first ?? "nothing" } // expected-note 20{{calls to instance method 'synchronous()' from outside of its actor context are implicitly asynchronous}}
   func asynchronous() async -> String {
     super.superState += 4
     return synchronous()
@@ -754,3 +754,30 @@ func testCrossActorProtocol<T: P>(t: T) async {
   t.f() // expected-error{{call is 'async' but is not marked with 'await'}}
   t.g() // expected-error{{call is 'async' but is not marked with 'await'}}
 }
+
+// ----------------------------------------------------------------------
+// @_
+// ----------------------------------------------------------------------
+func acceptAsyncSendableClosure<T>(_: @Sendable () async -> T) { }
+func acceptAsyncSendableClosureInheriting<T>(@_inheritActorContext _: @Sendable () async -> T) { }
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+extension MyActor {
+  func testSendableAndInheriting() {
+    acceptAsyncSendableClosure {
+      synchronous() // expected-error{{actor-isolated instance method 'synchronous()' cannot be referenced from a concurrent closure}}
+    }
+
+    acceptAsyncSendableClosure {
+      await synchronous() // ok
+    }
+
+    acceptAsyncSendableClosureInheriting {
+      synchronous() // okay
+    }
+
+    acceptAsyncSendableClosureInheriting {
+      await synchronous() // expected-warning{{no 'async' operations occur within 'await' expression}}
+    }
+  }
+}

test/Concurrency/global_actor_inference.swift

@@ -506,10 +506,34 @@ func acceptClosure<T>(_: () -> T) { }
 // ----------------------------------------------------------------------
 func takesUnsafeMainActor(@_unsafeMainActor fn: () -> Void) { }
 
-@MainActor func onlyOnMainActor() { }
+@MainActor func onlyOnMainActor() { } // expected-note{{calls to global function 'onlyOnMainActor()' from outside of its actor context are implicitly asynchronous}}
 
 func useUnsafeMainActor() {
   takesUnsafeMainActor {
     onlyOnMainActor() // okay due to parameter attribute
   }
 }
+
+// ----------------------------------------------------------------------
+// @_inheritActorContext
+// ----------------------------------------------------------------------
+func acceptAsyncSendableClosure<T>(_: @Sendable () async -> T) { }
+func acceptAsyncSendableClosureInheriting<T>(@_inheritActorContext _: @Sendable () async -> T) { }
+
+@MainActor func testCallFromMainActor() {
+  acceptAsyncSendableClosure {
+    onlyOnMainActor() // expected-error{{call to main actor-isolated global function 'onlyOnMainActor()' in a synchronous nonisolated context}}
+  }
+
+  acceptAsyncSendableClosure {
+    await onlyOnMainActor() // okay
+  }
+
+  acceptAsyncSendableClosureInheriting {
+    onlyOnMainActor() // okay
+  }
+
+  acceptAsyncSendableClosureInheriting {
+    await onlyOnMainActor() // expected-warning{{no 'async' operations occur within 'await' expression}}
+  }
+}