Merge pull request #71143 from hborla/isolation-inheritance

hborla · web-flow · commit 7bcade6cd3ec · 2024-01-25T10:11:50.000-08:00
[Concurrency] More precise modeling of `ActorIsolation` for isolated arguments and captures.
diff --git a/include/swift/AST/ActorIsolation.h b/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),
@@ -104,17 +106,23 @@ class ActorIsolation {
     return ActorIsolation(unsafe ? NonisolatedUnsafe : Nonisolated, nullptr);
   }
 
-  static ActorIsolation forActorInstanceSelf(NominalTypeDecl *actor) {
-    return ActorIsolation(ActorInstance, actor, 0);
-  }
+  static ActorIsolation forActorInstanceSelf(ValueDecl *decl);
 
   static ActorIsolation forActorInstanceParameter(NominalTypeDecl *actor,
                                                   unsigned parameterIndex) {
     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);
+
   static ActorIsolation forActorInstanceCapture(VarDecl *capturedActor) {
-    return ActorIsolation(ActorInstance, capturedActor);
+    return ActorIsolation(ActorInstance, capturedActor, 0);
   }
 
   static ActorIsolation forGlobalActor(Type globalActor) {
@@ -179,6 +187,8 @@ class ActorIsolation {
 
   VarDecl *getActorInstance() const;
 
+  Expr *getActorInstanceExpr() const;
+
   bool isGlobalActor() const {
     return getKind() == GlobalActor;
   }
@@ -213,27 +223,12 @@ class ActorIsolation {
   /// Substitute into types within the actor isolation.
   ActorIsolation subst(SubstitutionMap subs) const;
 
+  static bool isEqual(const ActorIsolation &lhs,
+               const ActorIsolation &rhs);
+
   friend bool operator==(const ActorIsolation &lhs,
                          const ActorIsolation &rhs) {
-    if (lhs.isGlobalActor() && rhs.isGlobalActor())
-      return areTypesEqual(lhs.globalActor, rhs.globalActor);
-
-    if (lhs.getKind() != rhs.getKind())
-      return false;
-
-    switch (lhs.getKind()) {
-    case Nonisolated:
-    case NonisolatedUnsafe:
-    case Unspecified:
-      return true;
-
-    case ActorInstance:
-      return (lhs.getActor() == rhs.getActor() &&
-              lhs.parameterIndex == rhs.parameterIndex);
-
-    case GlobalActor:
-      llvm_unreachable("Global actors handled above");
-    }
+    return ActorIsolation::isEqual(lhs, rhs);
   }
 
   friend bool operator!=(const ActorIsolation &lhs,
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -11271,18 +11271,82 @@ 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) {}
+
+ActorIsolation
+ActorIsolation::forActorInstanceParameter(Expr *actor,
+                                          unsigned parameterIndex) {
+  auto &ctx = actor->getType()->getASTContext();
+
+  // An isolated value of `nil` is statically nonisolated.
+  // FIXME: Also allow 'Optional.none'
+  if (dyn_cast<NilLiteralExpr>(actor))
+    return ActorIsolation::forNonisolated(/*unsafe*/false);
+
+  // An isolated value of `<global actor type>.shared` is statically
+  // global actor isolated.
+  if (auto *memberRef = dyn_cast<MemberRefExpr>(actor)) {
+    // Check that the member declaration witnesses the `shared`
+    // requirement of the `GlobalActor` protocol.
+    auto declRef = memberRef->getDecl();
+    auto baseType =
+        memberRef->getBase()->getType()->getMetatypeInstanceType();
+    if (auto globalActor = ctx.getProtocol(KnownProtocolKind::GlobalActor)) {
+      auto *dc = declRef.getDecl()->getDeclContext();
+      auto *module = dc->getParentModule();
+      auto conformance = module->checkConformance(baseType, globalActor);
+      if (conformance &&
+          conformance.getWitnessByName(baseType, ctx.Id_shared) == declRef) {
+        return ActorIsolation::forGlobalActor(baseType);
+      }
+    }
+  }
+
+  return ActorIsolation(ActorInstance, actor, parameterIndex + 1);
+}
+
+ActorIsolation
+ActorIsolation::forActorInstanceSelf(ValueDecl *decl) {
+  if (auto *fn = dyn_cast<AbstractFunctionDecl>(decl))
+    return ActorIsolation(ActorInstance, fn->getImplicitSelfDecl(), 0);
+
+  if (auto *storage = dyn_cast<AbstractStorageDecl>(decl)) {
+    if (auto *fn = storage->getAccessor(AccessorKind::Get)) {
+      return ActorIsolation(ActorInstance, fn->getImplicitSelfDecl(), 0);
+    }
+  }
+
+  auto *dc = decl->getDeclContext();
+  return ActorIsolation(ActorInstance, dc->getSelfNominalTypeDecl(), 0);
+}
 
 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();
   }
 
@@ -11298,6 +11362,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;
@@ -11317,6 +11390,39 @@ bool ActorIsolation::isDistributedActor() const {
   return getKind() == ActorInstance && getActor()->isDistributedActor();
 }
 
+bool ActorIsolation::isEqual(const ActorIsolation &lhs,
+                             const ActorIsolation &rhs) {
+  if (lhs.getKind() != rhs.getKind())
+    return false;
+
+  switch (lhs.getKind()) {
+  case Nonisolated:
+  case NonisolatedUnsafe:
+  case Unspecified:
+    return true;
+
+  case ActorInstance: {
+    auto *lhsActor = lhs.getActorInstance();
+    auto *rhsActor = rhs.getActorInstance();
+    if (lhsActor && rhsActor) {
+      // FIXME: This won't work for arbitrary isolated parameter captures.
+      if ((lhsActor->isSelfParameter() && rhsActor->isSelfParamCapture()) ||
+          (lhsActor->isSelfParamCapture() && rhsActor->isSelfParameter())) {
+        return true;
+      }
+    }
+
+    // The parameter index doesn't matter; only the actor instance
+    // values must be equal.
+    return (lhs.getActor() == rhs.getActor() &&
+            lhs.actorInstance == rhs.actorInstance);
+  }
+
+  case GlobalActor:
+    return areTypesEqual(lhs.globalActor, rhs.globalActor);
+  }
+}
+
 BuiltinTupleDecl::BuiltinTupleDecl(Identifier Name, DeclContext *Parent)
     : NominalTypeDecl(DeclKind::BuiltinTuple, Parent, Name, SourceLoc(),
                       ArrayRef<InheritedEntry>(), nullptr) {}
diff --git a/lib/AST/Type.cpp b/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();
diff --git a/lib/Sema/TypeCheckAttr.cpp b/lib/Sema/TypeCheckAttr.cpp
@@ -6832,7 +6832,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);
   }
 }
diff --git a/lib/Sema/TypeCheckConcurrency.cpp b/lib/Sema/TypeCheckConcurrency.cpp
@@ -141,7 +141,8 @@ bool swift::usesFlowSensitiveIsolation(AbstractFunctionDecl const *fn) {
       return true;
 
     // construct an isolation corresponding to the type.
-    auto actorTypeIso = ActorIsolation::forActorInstanceSelf(nominal);
+    auto actorTypeIso = ActorIsolation::forActorInstanceSelf(
+        const_cast<AbstractFunctionDecl *>(fn));
 
     return requiresFlowIsolation(actorTypeIso, cast<ConstructorDecl>(fn));
   }
@@ -3280,6 +3281,17 @@ namespace {
           continue;
 
         auto *arg = args->getExpr(paramIdx);
+
+        // FIXME: CurrentContextIsolationExpr does not have its actor set
+        // at this point.
+        if (auto *macro = dyn_cast<MacroExpansionExpr>(arg)) {
+          auto *expansion = macro->getRewritten();
+          if (auto *isolation = dyn_cast<CurrentContextIsolationExpr>(expansion)) {
+            recordCurrentContextIsolation(isolation);
+            arg = isolation->getActor();
+          }
+        }
+
         argForIsolatedParam = arg;
         if (getIsolatedActor(arg))
           continue;
@@ -3290,17 +3302,20 @@ 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);
         }
 
-        // FIXME: Also allow 'Optional.none'
-        if (dyn_cast<NilLiteralExpr>(arg->findOriginalValue())) {
-          mayExitToNonisolated = true;
-        } else {
-          unsatisfiedIsolation =
-              ActorIsolation::forActorInstanceParameter(nominal, paramIdx);
-          mayExitToNonisolated = false;
+        auto calleeIsolation = ActorIsolation::forActorInstanceParameter(
+            const_cast<Expr *>(arg->findOriginalValue()), paramIdx);
+
+        if (getContextIsolation() != calleeIsolation) {
+          if (calleeIsolation.isNonisolated()) {
+            mayExitToNonisolated = true;
+          } else {
+            unsatisfiedIsolation = calleeIsolation;
+          }
         }
 
         if (!fnType->getExtInfo().isAsync())
@@ -3447,6 +3462,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();
@@ -4661,7 +4681,7 @@ ActorIsolation ActorIsolationRequest::evaluate(
   if (evaluateOrDefault(evaluator, HasIsolatedSelfRequest{value}, false)) {
     auto actor = value->getDeclContext()->getSelfNominalTypeDecl();
     assert(actor && "could not find the actor that 'self' is isolated to");
-    return ActorIsolation::forActorInstanceSelf(actor);
+    return ActorIsolation::forActorInstanceSelf(value);
   }
 
   // If this declaration has an isolated parameter, it's isolated to that
@@ -4698,7 +4718,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
@@ -6092,7 +6112,7 @@ static ActorIsolation getActorIsolationForReference(
     // as needing to enter the actor.
     if (auto nominal = ctor->getDeclContext()->getSelfNominalTypeDecl()) {
       if (nominal->isAnyActor())
-        return ActorIsolation::forActorInstanceSelf(nominal);
+        return ActorIsolation::forActorInstanceSelf(decl);
     }
 
     // Fall through to treat initializers like any other declaration.
@@ -6108,7 +6128,7 @@ static ActorIsolation getActorIsolationForReference(
         declIsolation.isNonisolated()) {
       if (auto nominal = var->getDeclContext()->getSelfNominalTypeDecl()) {
         if (nominal->isAnyActor())
-          return ActorIsolation::forActorInstanceSelf(nominal);
+          return ActorIsolation::forActorInstanceSelf(decl);
 
         auto nominalIsolation = getActorIsolation(nominal);
         if (nominalIsolation.isGlobalActor())
diff --git a/stdlib/public/Distributed/DistributedActor.swift b/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
     }
diff --git a/test/Concurrency/isolated_parameters.swift b/test/Concurrency/isolated_parameters.swift
diff --git a/test/Concurrency/optional_isolated_parameters.swift b/test/Concurrency/optional_isolated_parameters.swift

Original file line number	Diff line number	Diff line change
`@@ -6832,7 +6832,14 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {`
`6832`	`6832`	`}`
`6833`	`6833`
`6834`	`6834`	`void AttributeChecker::visitKnownToBeLocalAttr(KnownToBeLocalAttr *attr) {`
`6835`		`- if (!D->isImplicit()) {`
	`6835`	`+ auto &ctx = D->getASTContext();`
	`6836`	`+ auto *module = D->getDeclContext()->getParentModule();`
	`6837`	`+ auto *distributed = ctx.getLoadedModule(ctx.Id_Distributed);`
	`6838`	`+`
	`6839`	+ // FIXME: An explicit `_local` is used in the implementation of
	`6840`	+ // `DistributedActor.whenLocal`, which otherwise violates actor
	`6841`	`+ // isolation checking.`
	`6842`	`+ if (!D->isImplicit() && (module != distributed)) {`
`6836`	`6843`	`diagnoseAndRemoveAttr(attr, diag::distributed_local_cannot_be_used);`
`6837`	`6844`	`}`
`6838`	`6845`	`}`