[Concurrency] Model actor isolation to an arbitrary actor instance using

a VarDecl or Expr.

This generalization exposed a bug where distributed actor isolation checking
was skipped in some cases, including for the isolated call in `whenLocal`.
The `whenLocal` implementation violated distributed actor isolation because
despite the `__isLocal` dynamic check, the `self` value passed to the `body`
function argument is still not statically local. To workaround this, I
applied the `_local` modifier explicitly to `self` before the call, which
also necessitated allowing `_local` to be written explicitly in the Distributed
library.

Author: hborla

include/swift/AST/ActorIsolation.h

@@ -70,7 +70,7 @@ class ActorIsolation {
 
 private:
   union {
-    llvm::PointerUnion<NominalTypeDecl *, VarDecl *> actorInstance;
+    llvm::PointerUnion<NominalTypeDecl *, VarDecl *, Expr *> actorInstance;
     Type globalActor;
     void *pointer;
   };
@@ -84,7 +84,9 @@ class ActorIsolation {
 
   ActorIsolation(Kind kind, NominalTypeDecl *actor, unsigned parameterIndex);
 
-  ActorIsolation(Kind kind, VarDecl *capturedActor);
+  ActorIsolation(Kind kind, VarDecl *actor, unsigned parameterIndex);
+
+  ActorIsolation(Kind kind, Expr *actor, unsigned parameterIndex);
 
   ActorIsolation(Kind kind, Type globalActor)
       : globalActor(globalActor), kind(kind), isolatedByPreconcurrency(false),
@@ -113,8 +115,18 @@ class ActorIsolation {
     return ActorIsolation(ActorInstance, actor, parameterIndex + 1);
   }
 
+  static ActorIsolation forActorInstanceParameter(VarDecl *actor,
+                                                  unsigned parameterIndex) {
+    return ActorIsolation(ActorInstance, actor, parameterIndex + 1);
+  }
+
+  static ActorIsolation forActorInstanceParameter(Expr *actor,
+                                                  unsigned parameterIndex) {
+    return ActorIsolation(ActorInstance, actor, parameterIndex + 1);
+  }
+
   static ActorIsolation forActorInstanceCapture(VarDecl *capturedActor) {
-    return ActorIsolation(ActorInstance, capturedActor);
+    return ActorIsolation(ActorInstance, capturedActor, 0);
   }
 
   static ActorIsolation forGlobalActor(Type globalActor) {
@@ -179,6 +191,8 @@ class ActorIsolation {
 
   VarDecl *getActorInstance() const;
 
+  Expr *getActorInstanceExpr() const;
+
   bool isGlobalActor() const {
     return getKind() == GlobalActor;
   }

lib/AST/Decl.cpp

@@ -11223,18 +11223,35 @@ ActorIsolation::ActorIsolation(Kind kind, NominalTypeDecl *actor,
     : actorInstance(actor), kind(kind), isolatedByPreconcurrency(false),
       silParsed(false), parameterIndex(parameterIndex) {}
 
-ActorIsolation::ActorIsolation(Kind kind, VarDecl *capturedActor)
-    : actorInstance(capturedActor), kind(kind), isolatedByPreconcurrency(false),
-      silParsed(false), parameterIndex(0) {}
+ActorIsolation::ActorIsolation(Kind kind, VarDecl *actor,
+                               unsigned parameterIndex)
+    : actorInstance(actor), kind(kind), isolatedByPreconcurrency(false),
+      silParsed(false), parameterIndex(parameterIndex) {}
+
+ActorIsolation::ActorIsolation(Kind kind, Expr *actor,
+                               unsigned parameterIndex)
+    : actorInstance(actor), kind(kind), isolatedByPreconcurrency(false),
+      silParsed(false), parameterIndex(parameterIndex) {}
 
 NominalTypeDecl *ActorIsolation::getActor() const {
   assert(getKind() == ActorInstance);
 
   if (silParsed)
     return nullptr;
 
+  Type actorType;
+
   if (auto *instance = actorInstance.dyn_cast<VarDecl *>()) {
-    return instance->getTypeInContext()
+    actorType = instance->getTypeInContext();
+  } else if (auto *instance = actorInstance.dyn_cast<Expr *>()) {
+    actorType = instance->getType();
+  }
+
+  if (actorType) {
+    if (auto wrapped = actorType->getOptionalObjectType()) {
+      actorType = wrapped;
+    }
+    return actorType
         ->getReferenceStorageReferent()->getAnyActor();
   }
 
@@ -11250,6 +11267,15 @@ VarDecl *ActorIsolation::getActorInstance() const {
   return actorInstance.dyn_cast<VarDecl *>();
 }
 
+Expr *ActorIsolation::getActorInstanceExpr() const {
+  assert(getKind() == ActorInstance);
+
+  if (silParsed)
+    return nullptr;
+
+  return actorInstance.dyn_cast<Expr *>();
+}
+
 bool ActorIsolation::isMainActor() const {
   if (silParsed)
     return false;

lib/AST/Type.cpp

@@ -549,6 +549,10 @@ NominalTypeDecl *TypeBase::getAnyActor() {
     return nullptr;
   }
 
+  if (auto self = getAs<DynamicSelfType>()) {
+    return self->getSelfType()->getAnyActor();
+  }
+
   // Existential types: check for Actor protocol.
   if (isExistentialType()) {
     auto layout = getExistentialLayout();

lib/Sema/TypeCheckAttr.cpp

@@ -6833,7 +6833,14 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {
 }
 
 void AttributeChecker::visitKnownToBeLocalAttr(KnownToBeLocalAttr *attr) {
-  if (!D->isImplicit()) {
+  auto &ctx = D->getASTContext();
+  auto *module = D->getDeclContext()->getParentModule();
+  auto *distributed = ctx.getLoadedModule(ctx.Id_Distributed);
+
+  // FIXME: An explicit `_local` is used in the implementation of
+  // `DistributedActor.whenLocal`, which otherwise violates actor
+  // isolation checking.
+  if (!D->isImplicit() && (module != distributed)) {
     diagnoseAndRemoveAttr(attr, diag::distributed_local_cannot_be_used);
   }
 }

lib/Sema/TypeCheckConcurrency.cpp

@@ -3290,16 +3290,19 @@ namespace {
         // that can talk about specific parameters.
         auto nominal = getType(arg)->getAnyNominal();
         if (!nominal) {
+          // FIXME: This is wrong for distributed actors.
           nominal = getType(arg)->getASTContext().getProtocol(
               KnownProtocolKind::Actor);
         }
 
+        auto origArg = const_cast<Expr *>(arg->findOriginalValue());
+
         // FIXME: Also allow 'Optional.none'
-        if (dyn_cast<NilLiteralExpr>(arg->findOriginalValue())) {
+        if (dyn_cast<NilLiteralExpr>(origArg)) {
           mayExitToNonisolated = true;
         } else {
           unsatisfiedIsolation =
-              ActorIsolation::forActorInstanceParameter(nominal, paramIdx);
+              ActorIsolation::forActorInstanceParameter(origArg, paramIdx);
           mayExitToNonisolated = false;
         }
 
@@ -3447,6 +3450,11 @@ namespace {
       Type optionalAnyActorType = isolationExpr->getType();
       switch (isolation) {
       case ActorIsolation::ActorInstance: {
+        if (auto *instance = isolation.getActorInstanceExpr()) {
+          actorExpr = instance;
+          break;
+        }
+
         const VarDecl *var = isolation.getActorInstance();
         if (!var) {
           auto dc = getDeclContext();
@@ -4690,7 +4698,7 @@ ActorIsolation ActorIsolationRequest::evaluate(
       actorType = paramType;
     }
     if (auto actor = actorType->getAnyActor())
-      return ActorIsolation::forActorInstanceParameter(actor, *paramIdx);
+      return ActorIsolation::forActorInstanceParameter(param, *paramIdx);
   }
 
   // Diagnose global state that is not either immutable plus Sendable or

stdlib/public/Distributed/DistributedActor.swift

@@ -366,7 +366,8 @@ extension DistributedActor {
     _ body: @Sendable (isolated Self) async throws -> T
   ) async rethrows -> T? {
     if __isLocalActor(self) {
-       return try await body(self)
+       _local let localSelf = self
+       return try await body(localSelf)
     } else {
       return nil
     }