[Concurrency] Annotate closures with their actor isolation.

Compute the actor isolation for every closure, noting whether it is
part of an actor instance (and which 'self' variable describes the
instance), global actor, or independent of any actor. This information
is required for propertly generating `hop_to_executor` instructions in
SIL.

Fixes rdar://71126554.

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/TypeCheckConcurrency.cpp

@@ -608,6 +608,7 @@ namespace {
 
     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 };
       }
@@ -715,11 +716,8 @@ namespace {
       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.
@@ -1016,6 +1014,111 @@ namespace {
       }
       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();
+      }
+    }
+    }
   };
 }
 

test/Concurrency/closure_isolation.swift

@@ -0,0 +1,70 @@
+// RUN: %target-swift-frontend -dump-ast %s -enable-experimental-concurrency | %FileCheck %s
+// REQUIRES: concurrency
+
+func acceptClosure<T>(_: () -> T) { }
+func acceptEscapingClosure<T>(_: @escaping () -> T) { }
+
+func acceptAsyncClosure<T>(_: () async -> T) { }
+func acceptEscapingAsyncClosure<T>(_: @escaping () async -> T) { }
+
+actor class MyActor {
+  func method() async -> String { "" }
+}
+
+extension MyActor {
+  // CHECK-LABEL: testClosureIsolation
+  func testClosureIsolation() async {
+    // CHECK: acceptAsyncClosure
+    // CHECK: closure_expr
+    // CHECK-SAME: actor-isolated=closure_isolation.(file).MyActor extension.testClosureIsolation().self
+    acceptAsyncClosure { await method() }
+
+    // CHECK: acceptAsyncClosure
+    // CHECK: closure_expr
+    // CHECK-NOT: actor-isolated
+    acceptAsyncClosure { () async in print("hello") }
+
+    // CHECK: acceptEscapingAsyncClosure
+    // CHECK: closure_expr
+    // CHECK-NOT: actor-isolated
+    acceptEscapingAsyncClosure { await self.method() }
+
+    // CHECK: acceptEscapingAsyncClosure
+    // CHECK: closure_expr
+    // CHECK-NOT:actor-isolated
+    acceptEscapingAsyncClosure { () async in print("hello") }
+  }
+}
+
+actor class SomeActor { }
+
+@globalActor
+struct SomeGlobalActor {
+  static let shared = SomeActor()
+}
+
+func someAsyncFunc() async { }
+
+// CHECK-LABEL: someGlobalActorFunc
+@SomeGlobalActor func someGlobalActorFunc() async {
+  // CHECK: acceptAsyncClosure
+  // CHECK: closure_expr
+  // CHECK-SAME: global-actor-isolated=SomeGlobalActor
+  acceptAsyncClosure { await someAsyncFunc() }
+
+  // CHECK: acceptAsyncClosure
+  // CHECK: closure_expr
+  // CHECK-SAME: global-actor-isolated=SomeGlobalActor
+  acceptAsyncClosure { () async in print("hello") }
+
+  // CHECK: acceptEscapingAsyncClosure
+  // CHECK: closure_expr
+  // CHECK-NOT: actor-isolated
+  acceptEscapingAsyncClosure { await someAsyncFunc() }
+
+  // CHECK: acceptEscapingAsyncClosure
+  // CHECK: closure_expr
+  // CHECK-NOT:actor-isolated
+  acceptEscapingAsyncClosure { () async in print("hello") }
+
+}