Properly compute actor isolation of local functions with captures.

Non-Sendable local functions can be actor-isolated if they capture an
isolated parameter. Treat them as such. Fixes rdar://83090397.

Author: DougGregor

include/swift/AST/CaptureInfo.h

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

lib/AST/CaptureInfo.cpp

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

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);

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

@@ -639,6 +639,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())
+  }
 }