Consistently handle 'self' captures

Author: DougGregor

include/swift/AST/CaptureInfo.h

@@ -36,7 +36,7 @@ namespace swift {
 class ValueDecl;
 class FuncDecl;
 class OpaqueValueExpr;
-class ParamDecl;
+class VarDecl;
 
 /// CapturedValue includes both the declaration being captured, along with flags
 /// that indicate how it is captured.
@@ -220,8 +220,10 @@ class CaptureInfo {
     return StorageAndFlags.getPointer()->getOpaqueValue();
   }
 
-  /// Retrieve the isolated parameter that has been captured, if there is one.
-  ParamDecl *getIsolatedParamCapture() const;
+  /// Retrieve the variable corresponding to an isolated parameter that has
+  /// been captured, if there is one. This might be a capture variable
+  /// that was initialized with an isolated parameter.
+  VarDecl *getIsolatedParamCapture() const;
 
   SWIFT_DEBUG_DUMP;
   void print(raw_ostream &OS) const;

include/swift/AST/Decl.h

@@ -5135,6 +5135,9 @@ class VarDecl : public AbstractStorageDecl {
       Bits.VarDecl.IsSelfParamCapture = IsSelfParamCapture;
   }
 
+  /// Check whether this capture of the self param is actor-isolated.
+  bool isSelfParamCaptureIsolated() const;
+
   /// Determines if this var has an initializer expression that should be
   /// exposed to clients.
   ///

lib/AST/CaptureInfo.cpp

@@ -78,7 +78,7 @@ getLocalCaptures(SmallVectorImpl<CapturedValue> &Result) const {
   }
 }
 
-ParamDecl *CaptureInfo::getIsolatedParamCapture() const {
+VarDecl *CaptureInfo::getIsolatedParamCapture() const {
   if (!hasLocalCaptures())
     return nullptr;
 
@@ -89,13 +89,18 @@ ParamDecl *CaptureInfo::getIsolatedParamCapture() const {
     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;
+    // If we captured an isolated parameter, return it.
+    if (auto param = dyn_cast_or_null<ParamDecl>(capture.getDecl())) {
+      // If we have captured an isolated parameter, return it.
+      if (param->isIsolated())
+        return param;
+    }
+
+    // If we captured 'self', check whether it is (still) isolated.
+    if (auto var = dyn_cast_or_null<VarDecl>(capture.getDecl())) {
+      if (var->isSelfParamCapture() && var->isSelfParamCaptureIsolated())
+        return var;
+    }
   }
 
   return nullptr;

lib/AST/Decl.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/ASTMangler.h"
+#include "swift/AST/CaptureInfo.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/ExistentialLayout.h"
@@ -8609,6 +8610,41 @@ void ClassDecl::setSuperclass(Type superclass) {
     true);
 }
 
+bool VarDecl::isSelfParamCaptureIsolated() const {
+  assert(isSelfParamCapture());
+
+  // Find the "self" parameter that we captured and determine whether
+  // it is potentially isolated.
+  for (auto dc = getDeclContext(); dc; dc = dc->getParent()) {
+    if (auto func = dyn_cast<AbstractFunctionDecl>(dc)) {
+      if (auto selfDecl = func->getImplicitSelfDecl()) {
+        return selfDecl->isIsolated();
+      }
+
+      if (auto capture = func->getCaptureInfo().getIsolatedParamCapture())
+        return capture->isSelfParameter() || capture->isSelfParamCapture();
+    }
+
+    if (auto closure = dyn_cast<AbstractClosureExpr>(dc)) {
+      switch (auto isolation = closure->getActorIsolation()) {
+      case ClosureActorIsolation::Independent:
+      case ClosureActorIsolation::GlobalActor:
+        return false;
+
+      case ClosureActorIsolation::ActorInstance:
+        auto isolatedVar = isolation.getActorInstance();
+        return isolatedVar->isSelfParameter() ||
+            isolatedVar-isSelfParamCapture();
+      }
+    }
+
+    if (dc->isModuleScopeContext() || dc->isTypeContext())
+      break;
+  }
+
+  return false;
+}
+
 ActorIsolation swift::getActorIsolation(ValueDecl *value) {
   auto &ctx = value->getASTContext();
   return evaluateOrDefault(
@@ -8638,7 +8674,7 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
 
     case ClosureActorIsolation::ActorInstance: {
       auto selfDecl = isolation.getActorInstance();
-      auto actorClass = selfDecl->getType()->getRValueType()
+      auto actorClass = selfDecl->getType()->getReferenceStorageReferent()
           ->getClassOrBoundGenericClass();
       // FIXME: Doesn't work properly with generics
       assert(actorClass && "Bad closure actor isolation?");

lib/Sema/TypeCheckConcurrency.cpp

@@ -4328,20 +4328,10 @@ bool swift::isPotentiallyIsolatedActor(
   if (auto param = dyn_cast<ParamDecl>(var))
     return isIsolated(param);
 
-  if (var->isSelfParamCapture()) {
-    // Find the "self" parameter that we captured and determine whether
-    // it is potentially isolated.
-    for (auto dc = var->getDeclContext(); dc; dc = dc->getParent()) {
-      if (auto func = dyn_cast<AbstractFunctionDecl>(dc)) {
-        if (auto selfDecl = func->getImplicitSelfDecl()) {
-          return selfDecl->isIsolated();
-        }
-      }
-
-      if (dc->isModuleScopeContext() || dc->isTypeContext())
-        break;
-    }
-  }
+  // If this is a captured 'self', check whether the original 'self' is
+  // isolated.
+  if (var->isSelfParamCapture())
+    return var->isSelfParamCaptureIsolated();
 
   return false;
 }

test/Concurrency/actor_isolation.swift

@@ -589,6 +589,35 @@ func f() {
   }
 }
 
+// ----------------------------------------------------------------------
+// Local function isolation restrictions
+// ----------------------------------------------------------------------
+@available(SwiftStdlib 5.1, *)
+actor AnActorWithClosures {
+  var counter: Int = 0 // expected-note 2 {{mutation of this property is only permitted within the actor}}
+  func exec() {
+    acceptEscapingClosure { [unowned self] in
+      self.counter += 1
+
+      acceptEscapingClosure {
+        self.counter += 1
+
+        acceptEscapingClosure { [self] in
+          self.counter += 1
+        }
+
+        acceptConcurrentClosure { [self] in
+          self.counter += 1 // expected-error{{actor-isolated property 'counter' can not be mutated from a Sendable closure}}
+
+          acceptEscapingClosure {
+            self.counter += 1 // expected-error{{actor-isolated property 'counter' can not be mutated from a non-isolated context}}
+          }
+        }
+      }
+    }
+  }
+}
+
 // ----------------------------------------------------------------------
 // Local function isolation restrictions
 // ----------------------------------------------------------------------