[Concurrency] Store ActorIsolation in AbstractClosureExpr instead of

ClosureActorIsolation.

Author: hborla

include/swift/AST/ActorIsolation.h

@@ -83,17 +83,17 @@ class ActorIsolation {
 
 private:
   union {
-    NominalTypeDecl *actor;
+    llvm::PointerUnion<NominalTypeDecl *, VarDecl *> actorInstance;
     Type globalActor;
     void *pointer;
   };
   unsigned kind : 3;
   unsigned isolatedByPreconcurrency : 1;
   unsigned parameterIndex : 28;
 
-  ActorIsolation(Kind kind, NominalTypeDecl *actor, unsigned parameterIndex)
-      : actor(actor), kind(kind), isolatedByPreconcurrency(false),
-        parameterIndex(parameterIndex) { }
+  ActorIsolation(Kind kind, NominalTypeDecl *actor, unsigned parameterIndex);
+
+  ActorIsolation(Kind kind, VarDecl *capturedActor);
 
   ActorIsolation(Kind kind, Type globalActor)
       : globalActor(globalActor), kind(kind), isolatedByPreconcurrency(false),
@@ -117,6 +117,10 @@ class ActorIsolation {
     return ActorIsolation(ActorInstance, actor, parameterIndex + 1);
   }
 
+  static ActorIsolation forActorInstanceCapture(VarDecl *capturedActor) {
+    return ActorIsolation(ActorInstance, capturedActor);
+  }
+
   static ActorIsolation forGlobalActor(Type globalActor, bool unsafe) {
     return ActorIsolation(
         unsafe ? GlobalActorUnsafe : GlobalActor, globalActor);
@@ -151,10 +155,9 @@ class ActorIsolation {
     }
   }
 
-  NominalTypeDecl *getActor() const {
-    assert(getKind() == ActorInstance);
-    return actor;
-  }
+  NominalTypeDecl *getActor() const;
+
+  VarDecl *getCapturedActor() const;
 
   bool isGlobalActor() const {
     return getKind() == GlobalActor || getKind() == GlobalActorUnsafe;
@@ -200,7 +203,8 @@ class ActorIsolation {
       return true;
 
     case ActorInstance:
-      return lhs.actor == rhs.actor && lhs.parameterIndex == rhs.parameterIndex;
+      return (lhs.getActor() == rhs.getActor() &&
+              lhs.parameterIndex == rhs.parameterIndex);
 
     case GlobalActor:
     case GlobalActorUnsafe:

include/swift/AST/Expr.h

@@ -3871,14 +3871,15 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   ParameterList *parameterList;
 
   /// Actor isolation of the closure.
-  ClosureActorIsolation actorIsolation;
+  ActorIsolation actorIsolation;
 
 public:
   AbstractClosureExpr(ExprKind Kind, Type FnType, bool Implicit,
                       DeclContext *Parent)
       : DeclContext(DeclContextKind::AbstractClosureExpr, Parent),
         Expr(Kind, Implicit, FnType),
-        parameterList(nullptr) {
+        parameterList(nullptr),
+        actorIsolation(ActorIsolation::forUnspecified()) {
     Bits.AbstractClosureExpr.Discriminator = InvalidDiscriminator;
   }
 
@@ -3951,14 +3952,37 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   /// returns nullptr if the closure doesn't have a body
   BraceStmt *getBody() const;
 
-  ClosureActorIsolation getClosureActorIsolation() const {
+  ActorIsolation getActorIsolation() const {
     return actorIsolation;
   }
 
-  void setActorIsolation(ClosureActorIsolation actorIsolation) {
+  void setActorIsolation(ActorIsolation actorIsolation) {
     this->actorIsolation = actorIsolation;
   }
 
+  ClosureActorIsolation getClosureActorIsolation() const {
+    bool preconcurrency = actorIsolation.preconcurrency();
+
+    switch (actorIsolation) {
+    case ActorIsolation::Unspecified:
+    case ActorIsolation::Nonisolated:
+      return ClosureActorIsolation::forNonisolated(preconcurrency);
+
+    case ActorIsolation::ActorInstance:
+      return ClosureActorIsolation::forActorInstance(
+          actorIsolation.getCapturedActor(), preconcurrency);
+
+    case ActorIsolation::GlobalActor:
+    case ActorIsolation::GlobalActorUnsafe:
+      return ClosureActorIsolation::forGlobalActor(
+          actorIsolation.getGlobalActor(), preconcurrency);
+    }
+  }
+
+  void setActorIsolation(ClosureActorIsolation closureIsolation) {
+    this->actorIsolation = closureIsolation.getActorIsolation();
+  }
+
   static bool classof(const Expr *E) {
     return E->getKind() >= ExprKind::First_AbstractClosureExpr &&
            E->getKind() <= ExprKind::Last_AbstractClosureExpr;

lib/AST/Decl.cpp

@@ -10597,6 +10597,33 @@ void swift::simple_display(llvm::raw_ostream &out, AnyFunctionRef fn) {
     out << "closure";
 }
 
+ActorIsolation::ActorIsolation(Kind kind, NominalTypeDecl *actor,
+                               unsigned parameterIndex)
+    : actorInstance(actor), kind(kind), 
+      isolatedByPreconcurrency(false),
+      parameterIndex(parameterIndex) { }
+
+ActorIsolation::ActorIsolation(Kind kind, VarDecl *capturedActor)
+    : actorInstance(capturedActor), kind(kind),
+      isolatedByPreconcurrency(false), 
+      parameterIndex(0) { }
+
+NominalTypeDecl *ActorIsolation::getActor() const {
+  assert(getKind() == ActorInstance);
+
+  if (auto *instance = actorInstance.dyn_cast<VarDecl *>()) {
+    return instance->getTypeInContext()
+        ->getReferenceStorageReferent()->getAnyActor();
+  }
+
+  return actorInstance.get<NominalTypeDecl *>();
+}
+
+VarDecl *ActorIsolation::getCapturedActor() const {
+  assert(getKind() == ActorInstance);
+  return actorInstance.dyn_cast<VarDecl *>();
+}
+
 bool ActorIsolation::isMainActor() const {
   if (isGlobalActor()) {
     if (auto *nominal = getGlobalActor()->getAnyNominal())

lib/AST/Expr.cpp

@@ -1895,8 +1895,8 @@ ActorIsolation ClosureActorIsolation::getActorIsolation() const {
         selfDecl->getTypeInContext()->getReferenceStorageReferent()->getAnyActor();
     assert(actor && "Bad closure actor isolation?");
     // FIXME: This could be a parameter... or a capture... hmmm.
-    return ActorIsolation::forActorInstanceSelf(actor).withPreconcurrency(
-        preconcurrency());
+    return ActorIsolation::forActorInstanceCapture(
+        getActorInstance()).withPreconcurrency(preconcurrency());
   }
   }
 }