Merge pull request swiftlang#34775 from DougGregor/concurrency-closure-isolation

[Concurrency] Annotate closures with their actor isolation.

Author: DougGregor

include/swift/AST/Expr.h

@@ -3553,6 +3553,79 @@ class SequenceExpr final : public Expr,
   }
 };
 
+/// Actor isolation for a closure.
+class ClosureActorIsolation {
+public:
+  enum Kind {
+    /// The closure is independent of any actor.
+    Independent,
+
+    /// The closure is tied to the actor instance described by the given
+    /// \c VarDecl*, which is the (captured) `self` of an actor.
+    ActorInstance,
+
+    /// The closure is tied to the global actor described by the given type.
+    GlobalActor,
+  };
+
+private:
+    /// The actor to which this closure is isolated.
+    ///
+    /// There are three possible states:
+    ///   - NULL: The closure is independent of any actor.
+    ///   - VarDecl*: The 'self' variable for the actor instance to which
+    ///     this closure is isolated. It will always have a type that conforms
+    ///     to the \c Actor protocol.
+    ///   - Type: The type of the global actor on which
+  llvm::PointerUnion<VarDecl *, Type> storage;
+
+  ClosureActorIsolation(VarDecl *selfDecl) : storage(selfDecl) { }
+  ClosureActorIsolation(Type globalActorType) : storage(globalActorType) { }
+
+public:
+  ClosureActorIsolation() : storage() { }
+
+  static ClosureActorIsolation forIndependent() {
+    return ClosureActorIsolation();
+  }
+
+  static ClosureActorIsolation forActorInstance(VarDecl *selfDecl) {
+    return ClosureActorIsolation(selfDecl);
+  }
+
+  static ClosureActorIsolation forGlobalActor(Type globalActorType) {
+    return ClosureActorIsolation(globalActorType);
+  }
+
+  /// Determine the kind of isolation.
+  Kind getKind() const {
+    if (storage.isNull())
+      return Kind::Independent;
+
+    if (storage.is<VarDecl *>())
+      return Kind::ActorInstance;
+
+    return Kind::GlobalActor;
+  }
+
+  /// Whether the closure is isolated at all.
+  explicit operator bool() const {
+    return getKind() != Kind::Independent;
+  }
+
+  /// Whether the closure is isolated at all.
+  operator Kind() const {
+    return getKind();
+  }
+
+  VarDecl *getActorInstance() const {
+    return storage.dyn_cast<VarDecl *>();
+  }
+
+  Type getGlobalActor() const {
+    return storage.dyn_cast<Type>();
+  }
+};
 
 /// A base class for closure expressions.
 class AbstractClosureExpr : public DeclContext, public Expr {
@@ -3561,6 +3634,9 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   /// The set of parameters.
   ParameterList *parameterList;
 
+  /// Actor isolation of the closure.
+  ClosureActorIsolation actorIsolation;
+
 public:
   AbstractClosureExpr(ExprKind Kind, Type FnType, bool Implicit,
                       unsigned Discriminator, DeclContext *Parent)
@@ -3625,6 +3701,12 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   /// Only valid when \c hasSingleExpressionBody() is true.
   Expr *getSingleExpressionBody() const;
 
+  ClosureActorIsolation getActorIsolation() const { return actorIsolation; }
+
+  void setActorIsolation(ClosureActorIsolation actorIsolation) {
+    this->actorIsolation = actorIsolation;
+  }
+
   static bool classof(const Expr *E) {
     return E->getKind() >= ExprKind::First_AbstractClosureExpr &&
            E->getKind() <= ExprKind::Last_AbstractClosureExpr;

lib/AST/ASTDumper.cpp

@@ -2505,6 +2505,22 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     printCommon(E, name);
     PrintWithColorRAII(OS, DiscriminatorColor)
       << " discriminator=" << E->getDiscriminator();
+
+    switch (auto isolation = E->getActorIsolation()) {
+    case ClosureActorIsolation::Independent:
+      break;
+
+    case ClosureActorIsolation::ActorInstance:
+      PrintWithColorRAII(OS, CapturesColor) << " actor-isolated="
+        << isolation.getActorInstance()->printRef();
+      break;
+
+    case ClosureActorIsolation::GlobalActor:
+      PrintWithColorRAII(OS, CapturesColor) << " global-actor-isolated="
+        << isolation.getGlobalActor().getString();
+      break;
+    }
+
     if (!E->getCaptureInfo().isTrivial()) {
       OS << " ";
       E->getCaptureInfo().print(PrintWithColorRAII(OS, CapturesColor).getOS());

lib/Sema/MiscDiagnostics.cpp

@@ -4612,7 +4612,6 @@ void swift::performSyntacticExprDiagnostics(const Expr *E,
   if (ctx.LangOpts.EnableObjCInterop)
     diagDeprecatedObjCSelectors(DC, E);
   diagnoseConstantArgumentRequirement(E, DC);
-  checkActorIsolation(E, DC);
 }
 
 void swift::performStmtDiagnostics(const Stmt *S, DeclContext *DC) {

lib/Sema/TypeCheckConcurrency.cpp

@@ -578,8 +578,10 @@ findMemberReference(Expr *expr) {
   return None;
 }
 
-void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
-  class ActorIsolationWalker : public ASTWalker {
+namespace {
+  /// Check for adherence to the actor isolation rules, emitting errors
+  /// when actor-isolated declarations are used in an unsafe manner.
+  class ActorIsolationChecker : public ASTWalker {
     ASTContext &ctx;
     SmallVector<const DeclContext *, 4> contextStack;
 
@@ -588,20 +590,25 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
     }
 
   public:
-    ActorIsolationWalker(const DeclContext *dc) : ctx(dc->getASTContext()) {
+    ActorIsolationChecker(const DeclContext *dc) : ctx(dc->getASTContext()) {
       contextStack.push_back(dc);
     }
 
-    bool shouldWalkIntoSeparatelyCheckedClosure(ClosureExpr *expr) override {
-      return false;
-    }
-
     bool shouldWalkCaptureInitializerExpressions() override { return true; }
 
     bool shouldWalkIntoTapExpression() override { return false; }
 
+    bool walkToDeclPre(Decl *D) override {
+      // Don't walk into functions; they'll be handled separately.
+      if (isa<AbstractFunctionDecl>(D))
+        return false;
+
+      return true;
+    }
+
     std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
       if (auto *closure = dyn_cast<AbstractClosureExpr>(expr)) {
+        closure->setActorIsolation(determineClosureIsolation(closure));
         contextStack.push_back(closure);
         return { true, expr };
       }
@@ -709,11 +716,8 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
       while (useContext != defContext) {
         // If we find an escaping closure, it can be run concurrently.
         if (auto closure = dyn_cast<AbstractClosureExpr>(useContext)) {
-          if (auto type = closure->getType()) {
-            if (auto fnType = type->getAs<AnyFunctionType>())
-              if (!fnType->isNoEscape())
-                return true;
-          }
+          if (isEscapingClosure(closure))
+            return true;
         }
 
         // If we find a local function, it can escape and be run concurrently.
@@ -1010,10 +1014,144 @@ void swift::checkActorIsolation(const Expr *expr, const DeclContext *dc) {
       }
       llvm_unreachable("unhandled actor isolation kind!");
     }
+
+    /// Determine whether this closure is escaping.
+    static bool isEscapingClosure(const AbstractClosureExpr *closure) {
+      if (auto type = closure->getType()) {
+        if (auto fnType = type->getAs<AnyFunctionType>())
+          return !fnType->isNoEscape();
+      }
+
+      return true;
+    }
+
+    /// Determine the isolation of a particular closure.
+    ///
+    /// This function assumes that enclosing closures have already had their
+    /// isolation checked.
+    ClosureActorIsolation determineClosureIsolation(
+        AbstractClosureExpr *closure) {
+      // An escaping closure is always actor-independent.
+      if (isEscapingClosure(closure))
+        return ClosureActorIsolation::forIndependent();
+
+      // A non-escaping closure gets its isolation from its context.
+      Optional<ActorIsolation> parentIsolation;
+      auto parentDC = closure->getParent();
+      switch (parentDC->getContextKind()) {
+      case DeclContextKind::AbstractClosureExpr: {
+        auto parentClosureIsolation = cast<AbstractClosureExpr>(parentDC)
+          ->getActorIsolation();
+        switch (parentClosureIsolation) {
+        case ClosureActorIsolation::Independent:
+          parentIsolation = ActorIsolation::forIndependent(
+              ActorIndependentKind::Safe);
+          break;
+
+        case ClosureActorIsolation::ActorInstance: {
+          auto selfDecl = parentClosureIsolation.getActorInstance();
+          auto actorClass = selfDecl->getType()->getRValueType()
+              ->getClassOrBoundGenericClass();
+          assert(actorClass && "Bad closure actor isolation?");
+          parentIsolation = ActorIsolation::forActorInstance(actorClass);
+          break;
+        }
+
+        case ClosureActorIsolation::GlobalActor:
+          parentIsolation = ActorIsolation::forGlobalActor(
+              parentClosureIsolation.getGlobalActor());
+          break;
+        }
+        break;
+      }
+
+      case DeclContextKind::AbstractFunctionDecl:
+      case DeclContextKind::SubscriptDecl:
+        parentIsolation = getActorIsolation(
+            cast<ValueDecl>(parentDC->getAsDecl()));
+        break;
+
+      case DeclContextKind::EnumElementDecl:
+      case DeclContextKind::ExtensionDecl:
+      case DeclContextKind::FileUnit:
+      case DeclContextKind::GenericTypeDecl:
+      case DeclContextKind::Initializer:
+      case DeclContextKind::Module:
+      case DeclContextKind::SerializedLocal:
+      case DeclContextKind::TopLevelCodeDecl:
+        return ClosureActorIsolation::forIndependent();
+      }
+
+      // We must have parent isolation determined to get here.
+      assert(parentIsolation && "Missing parent isolation?");
+      switch (*parentIsolation) {
+      case ActorIsolation::Independent:
+      case ActorIsolation::IndependentUnsafe:
+      case ActorIsolation::Unspecified:
+        return ClosureActorIsolation::forIndependent();
+
+      case ActorIsolation::GlobalActor: {
+        Type globalActorType = closure->mapTypeIntoContext(
+            parentIsolation->getGlobalActor()->mapTypeOutOfContext());
+        return ClosureActorIsolation::forGlobalActor(globalActorType);
+      }
+
+      case ActorIsolation::ActorInstance: {
+        SmallVector<CapturedValue, 2> localCaptures;
+        closure->getCaptureInfo().getLocalCaptures(localCaptures);
+        for (const auto &localCapture : localCaptures) {
+          if (localCapture.isDynamicSelfMetadata())
+            continue;
+
+          auto var = dyn_cast_or_null<VarDecl>(localCapture.getDecl());
+          if (!var)
+            continue;
+
+          // If we have captured the 'self' parameter, the closure is isolated
+          // to that actor instance.
+          if (var->isSelfParameter()) {
+            return ClosureActorIsolation::forActorInstance(var);
+          }
+        }
+
+        // When 'self' is not captured, this closure is actor-independent.
+        return ClosureActorIsolation::forIndependent();
+      }
+    }
+    }
   };
+}
+
+void swift::checkTopLevelActorIsolation(TopLevelCodeDecl *decl) {
+  ActorIsolationChecker checker(decl);
+  decl->getBody()->walk(checker);
+}
+
+void swift::checkFunctionActorIsolation(AbstractFunctionDecl *decl) {
+  ActorIsolationChecker checker(decl);
+  if (auto body = decl->getBody()) {
+    body->walk(checker);
+  }
+  if (auto ctor = dyn_cast<ConstructorDecl>(decl))
+    if (auto superInit = ctor->getSuperInitCall())
+      superInit->walk(checker);
+}
+
+void swift::checkInitializerActorIsolation(Initializer *init, Expr *expr) {
+  ActorIsolationChecker checker(init);
+  expr->walk(checker);
+}
+
+void swift::checkEnumElementActorIsolation(
+    EnumElementDecl *element, Expr *expr) {
+  ActorIsolationChecker checker(element);
+  expr->walk(checker);
+}
 
-  ActorIsolationWalker walker(dc);
-  const_cast<Expr *>(expr)->walk(walker);
+void swift::checkPropertyWrapperActorIsolation(
+   PatternBindingDecl *binding, Expr *expr) {
+  ActorIsolationChecker checker(binding->getDeclContext());
+  expr->walk(checker);
 }
 
 /// Determine actor isolation solely from attributes.

lib/Sema/TypeCheckConcurrency.h

@@ -22,14 +22,19 @@
 
 namespace swift {
 
+class AbstractFunctionDecl;
 class ActorIsolation;
 class ASTContext;
 class ClassDecl;
 class ConcreteDeclRef;
 class Decl;
 class DeclContext;
+class EnumElementDecl;
 class Expr;
 class FuncDecl;
+class Initializer;
+class PatternBindingDecl;
+class TopLevelCodeDecl;
 class TypeBase;
 class ValueDecl;
 
@@ -38,7 +43,12 @@ class ValueDecl;
 void addAsyncNotes(FuncDecl *func);
 
 /// Check actor isolation rules.
-void checkActorIsolation(const Expr *expr, const DeclContext *dc);
+void checkTopLevelActorIsolation(TopLevelCodeDecl *decl);
+void checkFunctionActorIsolation(AbstractFunctionDecl *decl);
+void checkInitializerActorIsolation(Initializer *init, Expr *expr);
+void checkEnumElementActorIsolation(EnumElementDecl *element, Expr *expr);
+void checkPropertyWrapperActorIsolation(
+    PatternBindingDecl *binding, Expr *expr);
 
 /// The isolation restriction in effect for a given declaration that is
 /// referenced from source.

lib/Sema/TypeCheckDecl.cpp

@@ -1201,6 +1201,7 @@ EnumRawValuesRequest::evaluate(Evaluator &eval, EnumDecl *ED,
       if (TypeChecker::typeCheckExpression(
               exprToCheck, ED,
               /*contextualInfo=*/{rawTy, CTP_EnumCaseRawValue})) {
+        checkEnumElementActorIsolation(elt, exprToCheck);
         TypeChecker::checkEnumElementEffects(elt, exprToCheck);
       }
     }