[Concurrency] Handle partial application of actor-isolated methods.

include/swift/AST/DiagnosticsSema.def

@@ -4149,6 +4149,9 @@ ERROR(actor_isolated_self_independent_context,none,
         "actor-isolated %0 %1 can not be referenced from an "
         "'@actorIndependent' context",
         (DescriptiveDeclKind, DeclName))
+ERROR(actor_isolated_partial_apply,none,
+        "actor-isolated %0 %1 can not be partially applied",
+        (DescriptiveDeclKind, DeclName))
 WARNING(concurrent_access_local,none,
         "local %0 %1 is unsafe to reference in code that may execute "
         "concurrently",

lib/AST/Expr.cpp

@@ -2007,17 +2007,37 @@ Expr *AutoClosureExpr::getSingleExpressionBody() const {
 }
 
 Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
+  auto maybeUnwrapOpenExistential = [](Expr *expr) {
+    if (auto *openExistential = dyn_cast<OpenExistentialExpr>(expr)) {
+      expr = openExistential->getSubExpr()->getSemanticsProvidingExpr();
+      if (auto *ICE = dyn_cast<ImplicitConversionExpr>(expr))
+        expr = ICE->getSyntacticSubExpr();
+    }
+
+    return expr;
+  };
+
+  auto maybeUnwrapOptionalEval = [](Expr *expr) {
+    if (auto optEval = dyn_cast<OptionalEvaluationExpr>(expr))
+      expr = optEval->getSubExpr();
+    if (auto inject = dyn_cast<InjectIntoOptionalExpr>(expr))
+      expr = inject->getSubExpr();
+    if (auto erasure = dyn_cast<ErasureExpr>(expr))
+      expr = erasure->getSubExpr();
+    if (auto bind = dyn_cast<BindOptionalExpr>(expr))
+      expr = bind->getSubExpr();
+    return expr;
+  };
+
   switch (getThunkKind()) {
   case AutoClosureExpr::Kind::None:
     break;
 
   case AutoClosureExpr::Kind::SingleCurryThunk: {
     auto *body = getSingleExpressionBody();
     body = body->getSemanticsProvidingExpr();
-
-    if (auto *openExistential = dyn_cast<OpenExistentialExpr>(body)) {
-      body = openExistential->getSubExpr()->getSemanticsProvidingExpr();
-    }
+    body = maybeUnwrapOpenExistential(body);
+    body = maybeUnwrapOptionalEval(body);
 
     if (auto *outerCall = dyn_cast<ApplyExpr>(body)) {
       return outerCall->getFn();
@@ -2034,18 +2054,12 @@ Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {
                AutoClosureExpr::Kind::SingleCurryThunk);
       auto *innerBody = innerClosure->getSingleExpressionBody();
       innerBody = innerBody->getSemanticsProvidingExpr();
-
-      if (auto *openExistential = dyn_cast<OpenExistentialExpr>(innerBody)) {
-        innerBody = openExistential->getSubExpr()->getSemanticsProvidingExpr();
-        if (auto *ICE = dyn_cast<ImplicitConversionExpr>(innerBody))
-          innerBody = ICE->getSyntacticSubExpr();
-      }
+      innerBody = maybeUnwrapOpenExistential(innerBody);
+      innerBody = maybeUnwrapOptionalEval(innerBody);
 
       if (auto *outerCall = dyn_cast<ApplyExpr>(innerBody)) {
         if (auto *innerCall = dyn_cast<ApplyExpr>(outerCall->getFn())) {
-          if (auto *declRef = dyn_cast<DeclRefExpr>(innerCall->getFn())) {
-            return declRef;
-          }
+          return innerCall->getFn();
         }
       }
     }

lib/Sema/CSApply.cpp

@@ -925,7 +925,6 @@ namespace {
 
       case ValueOwnership::Owned:
       case ValueOwnership::Shared:
-        // Ensure that the argument type matches up exactly.
         auto selfArgTy = ParenType::get(context,
                                         selfParam.getPlainType(),
                                         selfParam.getParameterFlags());

lib/Sema/TypeCheckConcurrency.cpp

@@ -338,6 +338,55 @@ class IsolationRestriction {
 
 }
 
+namespace {
+  /// Describes the important parts of a partial apply thunk.
+  struct PartialApplyThunkInfo {
+    Expr *base;
+    Expr *fn;
+    bool isEscaping;
+  };
+}
+
+/// Try to decompose a call that might be an invocation of a partial apply
+/// thunk.
+static Optional<PartialApplyThunkInfo> decomposePartialApplyThunk(
+    ApplyExpr *apply, Expr *parent) {
+  // Check for a call to the outer closure in the thunk.
+  auto outerAutoclosure = dyn_cast<AutoClosureExpr>(apply->getFn());
+  if (!outerAutoclosure ||
+      outerAutoclosure->getThunkKind()
+        != AutoClosureExpr::Kind::DoubleCurryThunk)
+    return None;
+
+  auto memberFn = outerAutoclosure->getUnwrappedCurryThunkExpr();
+  if (!memberFn)
+    return None;
+
+  // Determine whether the partial apply thunk was immediately converted to
+  // noescape.
+  bool isEscaping = true;
+  if (auto conversion = dyn_cast_or_null<FunctionConversionExpr>(parent)) {
+    auto fnType = conversion->getType()->getAs<FunctionType>();
+    isEscaping = fnType && !fnType->isNoEscape();
+  }
+
+  return PartialApplyThunkInfo{apply->getArg(), memberFn, isEscaping};
+}
+
+/// Find the immediate member reference in the given expression.
+static Optional<std::pair<ValueDecl *, SourceLoc>>
+findMemberReference(Expr *expr) {
+  if (auto declRef = dyn_cast<DeclRefExpr>(expr))
+    return std::make_pair(declRef->getDecl(), declRef->getLoc());
+
+  if (auto otherCtor = dyn_cast<OtherConstructorDeclRefExpr>(expr)) {
+    return std::make_pair(
+        static_cast<ValueDecl *>(otherCtor->getDecl()), otherCtor->getLoc());
+  }
+
+  return None;
+}
+
 void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
   class ActorIsolationWalker : public ASTWalker {
     ASTContext &ctx;
@@ -377,18 +426,27 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
         return { true, expr };
       }
 
-      if (auto call = dyn_cast<SelfApplyExpr>(expr)) {
-        Expr *fn = call->getFn()->getValueProvidingExpr();
-        if (auto declRef = dyn_cast<DeclRefExpr>(fn)) {
-          checkMemberReference(
-              call->getArg(), declRef->getDecl(), declRef->getLoc());
-          call->getArg()->walk(*this);
-          return { false, expr };
+      if (auto apply = dyn_cast<ApplyExpr>(expr)) {
+        // If this is a call to a partial apply thunk, decompose it to check it
+        // like based on the original written syntax, e.g., "self.method".
+        if (auto partialApply = decomposePartialApplyThunk(
+                apply, Parent.getAsExpr())) {
+          if (auto memberRef = findMemberReference(partialApply->fn)) {
+            checkMemberReference(
+                partialApply->base, memberRef->first, memberRef->second,
+                partialApply->isEscaping);
+
+            partialApply->base->walk(*this);
+            return { false, expr };
+          }
         }
+      }
 
-        if (auto otherCtor = dyn_cast<OtherConstructorDeclRefExpr>(fn)) {
+      if (auto call = dyn_cast<SelfApplyExpr>(expr)) {
+        Expr *fn = call->getFn()->getValueProvidingExpr();
+        if (auto memberRef = findMemberReference(fn)) {
           checkMemberReference(
-              call->getArg(), otherCtor->getDecl(), otherCtor->getLoc());
+              call->getArg(), memberRef->first, memberRef->second);
           call->getArg()->walk(*this);
           return { false, expr };
         }
@@ -557,7 +615,8 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
 
     /// Check a reference with the given base expression to the given member.
     bool checkMemberReference(
-        Expr *base, ValueDecl *member, SourceLoc memberLoc) {
+        Expr *base, ValueDecl *member, SourceLoc memberLoc,
+        bool isEscapingPartialApply = false) {
       if (!base || !member)
         return false;
 
@@ -566,6 +625,7 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
         return false;
 
       case IsolationRestriction::ActorSelf: {
+        // Must reference actor-isolated state on 'self'.
         auto selfVar = getSelfReference(base);
         if (!selfVar) {
           ctx.Diags.diagnose(
@@ -583,6 +643,18 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
                    cast<ValueDecl>(selfVar->getDeclContext()->getAsDecl()))) {
           case ActorIsolation::ActorInstance:
           case ActorIsolation::ActorPrivileged:
+            // An escaping partial application of something that is part of
+            // the actor's isolated state is never permitted.
+            if (isEscapingPartialApply) {
+              ctx.Diags.diagnose(
+                  memberLoc, diag::actor_isolated_partial_apply,
+                  member->getDescriptiveKind(),
+                  member->getName());
+              noteIsolatedActorMember(member);
+              return true;
+            }
+            break;
+
           case ActorIsolation::Unspecified:
             // Okay
             break;
@@ -595,7 +667,7 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
                 member->getDescriptiveKind(),
                 member->getName());
             noteIsolatedActorMember(member);
-            break;
+            return true;
         }
 
         // Check whether we are in a context that will not execute concurrently

lib/Sema/TypeCheckEffects.cpp

@@ -1514,7 +1514,21 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
     ContextScope scope(*this, Context::forClosure(E));
     scope.enterSubFunction();
     scope.resetCoverageForAutoclosureBody();
+
+    // Curry thunks aren't actually a call to the asynchronous function.
+    // Assume that async is covered in such contexts.
+    switch (E->getThunkKind()) {
+    case AutoClosureExpr::Kind::DoubleCurryThunk:
+    case AutoClosureExpr::Kind::SingleCurryThunk:
+      Flags.set(ContextFlags::IsAsyncCovered);
+      break;
+
+    case AutoClosureExpr::Kind::None:
+      break;
+    }
+
     E->getBody()->walk(*this);
+
     scope.preserveCoverageFromAutoclosureBody();
     return ShouldNotRecurse;
   }

test/Concurrency/actor_isolation.swift

@@ -7,13 +7,16 @@ func globalFunc() { }
 func acceptClosure<T>(_: () -> T) { }
 func acceptEscapingClosure<T>(_: @escaping () -> T) { }
 
+func acceptAsyncClosure<T>(_: () async -> T) { }
+func acceptEscapingAsyncClosure<T>(_: @escaping () async -> T) { }
+
 // ----------------------------------------------------------------------
 // Actor state isolation restrictions
 // ----------------------------------------------------------------------
 actor class MySuperActor {
   var superState: Int = 25
 
-  func superMethod() { }
+  func superMethod() { } // expected-note 2 {{only asynchronous methods can be used outside the actor instance; do you want to add 'async'?}}
   func superAsyncMethod() async { }
 
   subscript (index: Int) -> String { // expected-note{{subscript declared here}}
@@ -28,7 +31,7 @@ actor class MyActor: MySuperActor {
   class func synchronousClass() { }
   static func synchronousStatic() { }
 
-  func synchronous() -> String { text.first ?? "nothing" } // expected-note 6{{only asynchronous methods can be used outside the actor instance; do you want to add 'async'?}}
+  func synchronous() -> String { text.first ?? "nothing" } // expected-note 18{{only asynchronous methods can be used outside the actor instance; do you want to add 'async'?}}
   func asynchronous() async -> String { synchronous() }
 }
 
@@ -46,6 +49,7 @@ extension MyActor {
     // @actorIndependent
     _ = actorIndependentFunc(otherActor: self)
     _ = actorIndependentVar
+
     actorIndependentVar = 17
     _ = self.actorIndependentFunc(otherActor: self)
     _ = self.actorIndependentVar
@@ -60,6 +64,16 @@ extension MyActor {
     _ = immutableGlobal
     _ = mutableGlobal // expected-warning{{reference to var 'mutableGlobal' is not concurrency-safe because it involves shared mutable state}}
 
+    // Partial application
+    _ = synchronous  // expected-error{{actor-isolated instance method 'synchronous()' can not be referenced from an '@actorIndependent' context}}
+    _ = super.superMethod // expected-error{{actor-isolated instance method 'superMethod()' can not be referenced from an '@actorIndependent' context}}
+    acceptClosure(synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can not be referenced from an '@actorIndependent' context}}
+    acceptClosure(self.synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can not be referenced from an '@actorIndependent' context}}
+    acceptClosure(otherActor.synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can only be referenced on 'self'}}
+    acceptEscapingClosure(synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can not be referenced from an '@actorIndependent' context}}}}
+    acceptEscapingClosure(self.synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can not be referenced from an '@actorIndependent'}}
+    acceptEscapingClosure(otherActor.synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can only be referenced on 'self'}}
+
     return 5
   }
 
@@ -137,6 +151,19 @@ extension MyActor {
     }
 
     localVar = 0
+
+    // Partial application
+    _ = synchronous  // expected-error{{actor-isolated instance method 'synchronous()' can not be partially applied}}
+    _ = super.superMethod // expected-error{{actor-isolated instance method 'superMethod()' is unsafe to reference in code that may execute concurrently}}
+    acceptClosure(synchronous)
+    acceptClosure(self.synchronous)
+    acceptClosure(otherActor.synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can only be referenced on 'self'}}
+    acceptEscapingClosure(synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can not be partially applied}}
+    acceptEscapingClosure(self.synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can not be partially applied}}
+    acceptEscapingClosure(otherActor.synchronous) // expected-error{{actor-isolated instance method 'synchronous()' can only be referenced on 'self'}}
+
+    acceptAsyncClosure(self.asynchronous)
+    acceptEscapingAsyncClosure(self.asynchronous)
   }
 }