Properly compute actor isolation of local functions with captures.

include/swift/AST/CaptureInfo.h

@@ -36,6 +36,7 @@ namespace swift {
 class ValueDecl;
 class FuncDecl;
 class OpaqueValueExpr;
+class VarDecl;
 
 /// CapturedValue includes both the declaration being captured, along with flags
 /// that indicate how it is captured.
@@ -219,6 +220,11 @@ class CaptureInfo {
     return StorageAndFlags.getPointer()->getOpaqueValue();
   }
 
+  /// Retrieve the variable corresponding to an isolated parameter that has
+  /// been captured, if there is one. This might be a capture variable
+  /// that was initialized with an isolated parameter.
+  VarDecl *getIsolatedParamCapture() const;
+
   SWIFT_DEBUG_DUMP;
   void print(raw_ostream &OS) const;
 };

include/swift/AST/Decl.h

@@ -5135,6 +5135,9 @@ class VarDecl : public AbstractStorageDecl {
       Bits.VarDecl.IsSelfParamCapture = IsSelfParamCapture;
   }
 
+  /// Check whether this capture of the self param is actor-isolated.
+  bool isSelfParamCaptureIsolated() const;
+
   /// Determines if this var has an initializer expression that should be
   /// exposed to clients.
   ///

lib/AST/CaptureInfo.cpp

@@ -78,6 +78,34 @@ getLocalCaptures(SmallVectorImpl<CapturedValue> &Result) const {
   }
 }
 
+VarDecl *CaptureInfo::getIsolatedParamCapture() const {
+  if (!hasLocalCaptures())
+    return nullptr;
+
+  for (const auto &capture : getCaptures()) {
+    if (!capture.getDecl()->isLocalCapture())
+      continue;
+
+    if (capture.isDynamicSelfMetadata())
+      continue;
+
+    // If we captured an isolated parameter, return it.
+    if (auto param = dyn_cast_or_null<ParamDecl>(capture.getDecl())) {
+      // If we have captured an isolated parameter, return it.
+      if (param->isIsolated())
+        return param;
+    }
+
+    // If we captured 'self', check whether it is (still) isolated.
+    if (auto var = dyn_cast_or_null<VarDecl>(capture.getDecl())) {
+      if (var->isSelfParamCapture() && var->isSelfParamCaptureIsolated())
+        return var;
+    }
+  }
+
+  return nullptr;
+}
+
 LLVM_ATTRIBUTE_USED void CaptureInfo::dump() const {
   print(llvm::errs());
   llvm::errs() << '\n';

lib/AST/Decl.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/ASTMangler.h"
+#include "swift/AST/CaptureInfo.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/ExistentialLayout.h"
@@ -8609,6 +8610,41 @@ void ClassDecl::setSuperclass(Type superclass) {
     true);
 }
 
+bool VarDecl::isSelfParamCaptureIsolated() const {
+  assert(isSelfParamCapture());
+
+  // Find the "self" parameter that we captured and determine whether
+  // it is potentially isolated.
+  for (auto dc = getDeclContext(); dc; dc = dc->getParent()) {
+    if (auto func = dyn_cast<AbstractFunctionDecl>(dc)) {
+      if (auto selfDecl = func->getImplicitSelfDecl()) {
+        return selfDecl->isIsolated();
+      }
+
+      if (auto capture = func->getCaptureInfo().getIsolatedParamCapture())
+        return capture->isSelfParameter() || capture->isSelfParamCapture();
+    }
+
+    if (auto closure = dyn_cast<AbstractClosureExpr>(dc)) {
+      switch (auto isolation = closure->getActorIsolation()) {
+      case ClosureActorIsolation::Independent:
+      case ClosureActorIsolation::GlobalActor:
+        return false;
+
+      case ClosureActorIsolation::ActorInstance:
+        auto isolatedVar = isolation.getActorInstance();
+        return isolatedVar->isSelfParameter() ||
+            isolatedVar-isSelfParamCapture();
+      }
+    }
+
+    if (dc->isModuleScopeContext() || dc->isTypeContext())
+      break;
+  }
+
+  return false;
+}
+
 ActorIsolation swift::getActorIsolation(ValueDecl *value) {
   auto &ctx = value->getASTContext();
   return evaluateOrDefault(
@@ -8638,7 +8674,7 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
 
     case ClosureActorIsolation::ActorInstance: {
       auto selfDecl = isolation.getActorInstance();
-      auto actorClass = selfDecl->getType()->getRValueType()
+      auto actorClass = selfDecl->getType()->getReferenceStorageReferent()
           ->getClassOrBoundGenericClass();
       // FIXME: Doesn't work properly with generics
       assert(actorClass && "Bad closure actor isolation?");

lib/SILGen/SILGenProlog.cpp

@@ -558,6 +558,17 @@ void SILGenFunction::emitProlog(CaptureInfo captureInfo,
       !F.isAsync() &&
       !isInActorDestructor(FunctionDC);
 
+  // Local function to load the expected executor from a local actor
+  auto loadExpectedExecutorForLocalVar = [&](VarDecl *var) {
+    auto loc = RegularLocation::getAutoGeneratedLocation(F.getLocation());
+    Type actorType = var->getType();
+    RValue actorInstanceRV = emitRValueForDecl(
+        loc, var, actorType, AccessSemantics::Ordinary);
+    ManagedValue actorInstance =
+        std::move(actorInstanceRV).getScalarValue();
+    ExpectedExecutor = emitLoadActorExecutor(loc, actorInstance);
+  };
+
   if (auto *funcDecl =
         dyn_cast_or_null<AbstractFunctionDecl>(FunctionDC->getAsDecl())) {
     auto actorIsolation = getActorIsolation(funcDecl);
@@ -569,13 +580,7 @@ void SILGenFunction::emitProlog(CaptureInfo captureInfo,
       if (F.isAsync()) {
         for (auto param : *funcDecl->getParameters()) {
           if (param->isIsolated()) {
-            auto loc = RegularLocation::getAutoGeneratedLocation(F.getLocation());
-            Type actorType = param->getType();
-            RValue actorInstanceRV = emitRValueForDecl(
-                loc, param, actorType, AccessSemantics::Ordinary);
-            ManagedValue actorInstance =
-                std::move(actorInstanceRV).getScalarValue();
-            ExpectedExecutor = emitLoadActorExecutor(loc, actorInstance);
+            loadExpectedExecutorForLocalVar(param);
             break;
           }
         }
@@ -588,16 +593,21 @@ void SILGenFunction::emitProlog(CaptureInfo captureInfo,
     }
 
     case ActorIsolation::ActorInstance: {
-      assert(selfParam && "no self parameter for ActorInstance isolation");
       // Only produce an executor for actor-isolated functions that are async
       // or are local functions. The former require a hop, while the latter
       // are prone to dynamic data races in code that does not enforce Sendable
       // completely.
       if (F.isAsync() ||
           (wantDataRaceChecks && funcDecl->isLocalCapture())) {
-        auto loc = RegularLocation::getAutoGeneratedLocation(F.getLocation());
-        ManagedValue selfArg = ManagedValue::forUnmanaged(F.getSelfArgument());
-        ExpectedExecutor = emitLoadActorExecutor(loc, selfArg);
+        if (auto isolatedParam = funcDecl->getCaptureInfo()
+                .getIsolatedParamCapture()) {
+          loadExpectedExecutorForLocalVar(isolatedParam);
+        } else {
+          assert(selfParam && "no self parameter for ActorInstance isolation");
+          auto loc = RegularLocation::getAutoGeneratedLocation(F.getLocation());
+          ManagedValue selfArg = ManagedValue::forUnmanaged(F.getSelfArgument());
+          ExpectedExecutor = emitLoadActorExecutor(loc, selfArg);
+        }
       }
       break;
     }
@@ -619,14 +629,7 @@ void SILGenFunction::emitProlog(CaptureInfo captureInfo,
 
     case ClosureActorIsolation::ActorInstance: {
       if (wantExecutor) {
-        auto loc = RegularLocation::getAutoGeneratedLocation(F.getLocation());
-        auto actorDecl = actorIsolation.getActorInstance();
-        Type actorType = actorDecl->getType();
-        RValue actorInstanceRV = emitRValueForDecl(loc,
-          actorDecl, actorType, AccessSemantics::Ordinary);
-        ManagedValue actorInstance =
-            std::move(actorInstanceRV).getScalarValue();
-        ExpectedExecutor = emitLoadActorExecutor(loc, actorInstance);
+        loadExpectedExecutorForLocalVar(actorIsolation.getActorInstance());
       }
       break;
     }

lib/Sema/TypeCheckConcurrency.cpp

@@ -2640,7 +2640,7 @@ namespace {
       if (isSendableClosure(closure, /*forActorIsolation=*/true))
         return ClosureActorIsolation::forIndependent();
 
-      // A non-escaping closure gets its isolation from its context.
+      // A non-Sendable closure gets its isolation from its context.
       auto parentIsolation = getActorIsolationOfContext(closure->getParent());
 
       // We must have parent isolation determined to get here.
@@ -2658,25 +2658,9 @@ namespace {
 
       case ActorIsolation::ActorInstance:
       case ActorIsolation::DistributedActorInstance: {
-        SmallVector<CapturedValue, 2> localCaptures;
-        closure->getCaptureInfo().getLocalCaptures(localCaptures);
-        for (const auto &localCapture : localCaptures) {
-          if (localCapture.isDynamicSelfMetadata())
-            continue;
-
-          auto param = dyn_cast_or_null<ParamDecl>(localCapture.getDecl());
-          if (!param)
-            continue;
-
-          // If we have captured an isolated parameter, the closure is isolated
-          // to that actor instance.
-          if (param->isIsolated()) {
-            return ClosureActorIsolation::forActorInstance(param);
-          }
-        }
+        if (auto param = closure->getCaptureInfo().getIsolatedParamCapture())
+          return ClosureActorIsolation::forActorInstance(param);
 
-        // When no actor instance  is not captured, this closure is
-        // actor-independent.
         return ClosureActorIsolation::forIndependent();
       }
     }
@@ -3328,19 +3312,23 @@ ActorIsolation ActorIsolationRequest::evaluate(
     return inferred;
   };
 
-  // If this is a "defer" function body, inherit the global actor isolation
-  // from its context.
+  // If this is a local function, inherit the actor isolation from its
+  // context if it global or was captured.
   if (auto func = dyn_cast<FuncDecl>(value)) {
-    if (func->isDeferBody()) {
+    if (func->isLocalCapture() && !func->isSendable()) {
       switch (auto enclosingIsolation =
                   getActorIsolationOfContext(func->getDeclContext())) {
-      case ActorIsolation::ActorInstance:
-      case ActorIsolation::DistributedActorInstance:
       case ActorIsolation::Independent:
       case ActorIsolation::Unspecified:
         // Do nothing.
         break;
 
+      case ActorIsolation::ActorInstance:
+      case ActorIsolation::DistributedActorInstance:
+        if (auto param = func->getCaptureInfo().getIsolatedParamCapture())
+          return inferredIsolation(enclosingIsolation);
+        break;
+
       case ActorIsolation::GlobalActor:
       case ActorIsolation::GlobalActorUnsafe:
         return inferredIsolation(enclosingIsolation);
@@ -4340,20 +4328,10 @@ bool swift::isPotentiallyIsolatedActor(
   if (auto param = dyn_cast<ParamDecl>(var))
     return isIsolated(param);
 
-  if (var->isSelfParamCapture()) {
-    // Find the "self" parameter that we captured and determine whether
-    // it is potentially isolated.
-    for (auto dc = var->getDeclContext(); dc; dc = dc->getParent()) {
-      if (auto func = dyn_cast<AbstractFunctionDecl>(dc)) {
-        if (auto selfDecl = func->getImplicitSelfDecl()) {
-          return selfDecl->isIsolated();
-        }
-      }
-
-      if (dc->isModuleScopeContext() || dc->isTypeContext())
-        break;
-    }
-  }
+  // If this is a captured 'self', check whether the original 'self' is
+  // isolated.
+  if (var->isSelfParamCapture())
+    return var->isSelfParamCaptureIsolated();
 
   return false;
 }

test/Concurrency/actor_definite_init.swift

@@ -447,6 +447,6 @@ actor Rain {
 
     defer { _ = self.x }
 
-    defer { Task { await self.f() } }
+    defer { Task { self.f() } }
   }
 }

test/Concurrency/actor_definite_init_swift6.swift

@@ -445,6 +445,6 @@ actor Rain {
 
     defer { _ = self.x }
 
-    defer { Task { await self.f() } }
+    defer { Task { self.f() } }
   }
 }

test/Concurrency/actor_isolation.swift

@@ -589,6 +589,35 @@ func f() {
   }
 }
 
+// ----------------------------------------------------------------------
+// Local function isolation restrictions
+// ----------------------------------------------------------------------
+@available(SwiftStdlib 5.1, *)
+actor AnActorWithClosures {
+  var counter: Int = 0 // expected-note 2 {{mutation of this property is only permitted within the actor}}
+  func exec() {
+    acceptEscapingClosure { [unowned self] in
+      self.counter += 1
+
+      acceptEscapingClosure {
+        self.counter += 1
+
+        acceptEscapingClosure { [self] in
+          self.counter += 1
+        }
+
+        acceptConcurrentClosure { [self] in
+          self.counter += 1 // expected-error{{actor-isolated property 'counter' can not be mutated from a Sendable closure}}
+
+          acceptEscapingClosure {
+            self.counter += 1 // expected-error{{actor-isolated property 'counter' can not be mutated from a non-isolated context}}
+          }
+        }
+      }
+    }
+  }
+}
+
 // ----------------------------------------------------------------------
 // Local function isolation restrictions
 // ----------------------------------------------------------------------
@@ -639,6 +668,27 @@ func checkLocalFunctions() async {
   print(k)
 }
 
+@available(SwiftStdlib 5.1, *)
+actor LocalFunctionIsolatedActor {
+  func a() -> Bool { // expected-note{{calls to instance method 'a()' from outside of its actor context are implicitly asynchronous}}
+    return true
+  }
+
+  func b() -> Bool {
+    func c() -> Bool {
+      return true && a() // okay, c is isolated
+    }
+    return c()
+  }
+
+  func b2() -> Bool {
+    @Sendable func c() -> Bool {
+      return true && a() // expected-error{{actor-isolated instance method 'a()' can not be referenced from a non-isolated context}}
+    }
+    return c()
+  }
+}
+
 // ----------------------------------------------------------------------
 // Lazy properties with initializers referencing 'self'
 // ----------------------------------------------------------------------

test/SILGen/check_executor.swift

@@ -46,4 +46,19 @@ public actor MyActor {
       return 5
     }
   }
+
+  // CHECK-CANONICAL-LABEL: sil private [ossa] @$s4test7MyActorC0A13LocalFunctionyyF5localL_SiyF : $@convention(thin) (@guaranteed MyActor) -> Int
+  // CHECK-CANONICAL: [[CAPTURE:%.*]] = copy_value %0 : $MyActor
+  // CHECK-CANONICAL-NEXT: [[BORROWED_CAPTURE:%.*]] = begin_borrow [[CAPTURE]] : $MyActor
+  // CHECK-CANONICAL-NEXT: [[EXECUTOR:%.*]] = builtin "buildDefaultActorExecutorRef"<MyActor>([[BORROWED_CAPTURE]] : $MyActor) : $Builtin.Executor
+  // CHECK-CANONICAL-NEXT: [[EXECUTOR_DEP:%.*]] = mark_dependence [[EXECUTOR]] : $Builtin.Executor on [[BORROWED_CAPTURE]] : $MyActor
+  // CHECK-CANONICAL: [[CHECK_FN:%.*]] = function_ref @$ss22_checkExpectedExecutor14_filenameStart01_D6Length01_D7IsASCII5_line9_executoryBp_BwBi1_BwBetF : $@convention(thin) (Builtin.RawPointer, Builtin.Word, Builtin.Int1, Builtin.Word, Builtin.Executor) -> ()
+  // CHECK-CANONICAL-NEXT: apply [[CHECK_FN]]({{.*}}, [[EXECUTOR_DEP]])
+  public func testLocalFunction() {
+    func local() -> Int {
+      return counter
+    }
+
+    print(local())
+  }
 }