[Sema] Add custom functions to ActorIsolationChecker to determine expr type and closure actor isolation

include/swift/AST/ActorIsolation.h

@@ -30,6 +30,19 @@ class VarDecl;
 class NominalTypeDecl;
 class SubstitutionMap;
 class AbstractFunctionDecl;
+class AbstractClosureExpr;
+class ClosureActorIsolation;
+
+/// Trampoline for AbstractClosureExpr::getActorIsolation.
+ClosureActorIsolation
+__AbstractClosureExpr_getActorIsolation(AbstractClosureExpr *CE);
+
+/// Returns a function reference to \c __AbstractClosureExpr_getActorIsolation.
+/// This is needed so we can use it as a default argument for
+/// \c getActorIsolationOfContext without knowing the layout of
+/// \c ClosureActorIsolation.
+llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+_getRef__AbstractClosureExpr_getActorIsolation();
 
 /// Determine whether the given types are (canonically) equal, declared here
 /// to avoid having to include Types.h.
@@ -209,7 +222,15 @@ class ActorIsolation {
 ActorIsolation getActorIsolation(ValueDecl *value);
 
 /// Determine how the given declaration context is isolated.
-ActorIsolation getActorIsolationOfContext(DeclContext *dc);
+/// \p getClosureActorIsolation allows the specification of actor isolation for
+/// closures that haven't been saved been saved to the AST yet. This is useful
+/// for solver-based code completion which doesn't modify the AST but stores the
+/// actor isolation of closures in the constraint system solution.
+ActorIsolation getActorIsolationOfContext(
+    DeclContext *dc,
+    llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+        getClosureActorIsolation =
+            _getRef__AbstractClosureExpr_getActorIsolation());
 
 /// Check if both the value, and context are isolated to the same actor.
 bool isSameActorIsolated(ValueDecl *value, DeclContext *dc);

include/swift/AST/Expr.h

@@ -3658,6 +3658,8 @@ class ClosureActorIsolation {
   bool preconcurrency() const {
     return storage.getInt();
   }
+
+  ActorIsolation getActorIsolation() const;
 };
 
 /// A base class for closure expressions.

include/swift/IDE/CompletionLookup.h

@@ -126,6 +126,10 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   llvm::SmallDenseMap<const VarDecl *, Type> SolutionSpecificVarTypes;
 
   bool CanCurrDeclContextHandleAsync = false;
+  /// Actor isolations that were determined during constraint solving but that
+  /// haven't been saved to the AST.
+  llvm::DenseMap<AbstractClosureExpr *, ClosureActorIsolation>
+      ClosureActorIsolations;
   bool HaveDot = false;
   bool IsUnwrappedOptional = false;
   SourceLoc DotLoc;
@@ -262,6 +266,12 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     this->CanCurrDeclContextHandleAsync = CanCurrDeclContextHandleAsync;
   }
 
+  void setClosureActorIsolations(
+      llvm::DenseMap<AbstractClosureExpr *, ClosureActorIsolation>
+          ClosureActorIsolations) {
+    this->ClosureActorIsolations = ClosureActorIsolations;
+  }
+
   const ExpectedTypeContext *getExpectedTypeContext() const {
     return &expectedTypeContext;
   }

include/swift/Sema/IDETypeChecking.h

@@ -327,6 +327,12 @@ namespace swift {
   /// Just a proxy to swift::contextUsesConcurrencyFeatures() from lib/IDE code.
   bool completionContextUsesConcurrencyFeatures(const DeclContext *dc);
 
+  /// Determine the isolation of a particular closure.
+  ClosureActorIsolation determineClosureActorIsolation(
+      AbstractClosureExpr *closure, llvm::function_ref<Type(Expr *)> getType,
+      llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+          getClosureActorIsolation);
+
   /// If the capture list shadows any declarations using shorthand syntax, i.e.
   /// syntax that names both the newly declared variable and the referenced
   /// variable by the same identifier in the source text, i.e. `[foo]`, return

lib/AST/Decl.cpp

@@ -9211,7 +9211,10 @@ ActorIsolation swift::getActorIsolation(ValueDecl *value) {
       ActorIsolation::forUnspecified());
 }
 
-ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
+ActorIsolation swift::getActorIsolationOfContext(
+    DeclContext *dc,
+    llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+        getClosureActorIsolation) {
   auto dcToUse = dc;
   // Defer bodies share actor isolation of their enclosing context.
   if (auto FD = dyn_cast<FuncDecl>(dcToUse)) {
@@ -9226,27 +9229,7 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
     return getActorIsolation(var);
 
   if (auto *closure = dyn_cast<AbstractClosureExpr>(dcToUse)) {
-    switch (auto isolation = closure->getActorIsolation()) {
-    case ClosureActorIsolation::Independent:
-      return ActorIsolation::forIndependent()
-                .withPreconcurrency(isolation.preconcurrency());
-
-    case ClosureActorIsolation::GlobalActor: {
-      return ActorIsolation::forGlobalActor(
-          isolation.getGlobalActor(), /*unsafe=*/false)
-                .withPreconcurrency(isolation.preconcurrency());
-    }
-
-    case ClosureActorIsolation::ActorInstance: {
-      auto selfDecl = isolation.getActorInstance();
-      auto actor = selfDecl->getType()->getReferenceStorageReferent()
-          ->getAnyActor();
-      assert(actor && "Bad closure actor isolation?");
-      // FIXME: This could be a parameter... or a capture... hmmm.
-      return ActorIsolation::forActorInstanceSelf(actor)
-        .withPreconcurrency(isolation.preconcurrency());
-    }
-    }
+    return getClosureActorIsolation(closure).getActorIsolation();
   }
 
   if (auto *tld = dyn_cast<TopLevelCodeDecl>(dcToUse)) {

lib/AST/Expr.cpp

@@ -1779,6 +1779,29 @@ RebindSelfInConstructorExpr::getCalledConstructor(bool &isChainToSuper) const {
   return otherCtorRef;
 }
 
+ActorIsolation ClosureActorIsolation::getActorIsolation() const {
+  switch (getKind()) {
+  case ClosureActorIsolation::Independent:
+    return ActorIsolation::forIndependent().withPreconcurrency(
+        preconcurrency());
+
+  case ClosureActorIsolation::GlobalActor: {
+    return ActorIsolation::forGlobalActor(getGlobalActor(), /*unsafe=*/false)
+        .withPreconcurrency(preconcurrency());
+  }
+
+  case ClosureActorIsolation::ActorInstance: {
+    auto selfDecl = getActorInstance();
+    auto actor =
+        selfDecl->getType()->getReferenceStorageReferent()->getAnyActor();
+    assert(actor && "Bad closure actor isolation?");
+    // FIXME: This could be a parameter... or a capture... hmmm.
+    return ActorIsolation::forActorInstanceSelf(actor).withPreconcurrency(
+        preconcurrency());
+  }
+  }
+}
+
 void AbstractClosureExpr::setParameterList(ParameterList *P) {
   parameterList = P;
   // Change the DeclContext of any parameters to be this closure.
@@ -1880,6 +1903,16 @@ Expr *AbstractClosureExpr::getSingleExpressionBody() const {
   return nullptr;
 }
 
+ClosureActorIsolation
+swift::__AbstractClosureExpr_getActorIsolation(AbstractClosureExpr *CE) {
+  return CE->getActorIsolation();
+}
+
+llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+swift::_getRef__AbstractClosureExpr_getActorIsolation() {
+  return __AbstractClosureExpr_getActorIsolation;
+}
+
 #define FORWARD_SOURCE_LOCS_TO(CLASS, NODE) \
   SourceRange CLASS::getSourceRange() const {     \
     return (NODE)->getSourceRange();              \

lib/IDE/CompletionLookup.cpp

@@ -755,8 +755,18 @@ void CompletionLookup::analyzeActorIsolation(
     }
     LLVM_FALLTHROUGH;
   case ActorIsolation::GlobalActor: {
-    auto contextIsolation =
-        getActorIsolationOfContext(const_cast<DeclContext *>(CurrDeclContext));
+    auto getClosureActorIsolation = [this](AbstractClosureExpr *CE) {
+      // Prefer solution-specific actor-isolations and fall back to the one
+      // recorded in the AST.
+      auto isolation = ClosureActorIsolations.find(CE);
+      if (isolation != ClosureActorIsolations.end()) {
+        return isolation->second;
+      } else {
+        return CE->getActorIsolation();
+      }
+    };
+    auto contextIsolation = getActorIsolationOfContext(
+        const_cast<DeclContext *>(CurrDeclContext), getClosureActorIsolation);
     if (contextIsolation != isolation) {
       implicitlyAsync = true;
     }

lib/SILGen/SILGenProlog.cpp

@@ -921,7 +921,10 @@ void SILGenFunction::emitHopToActorValue(SILLocation loc, ManagedValue actor) {
   if (!F.isAsync()) {
     llvm::report_fatal_error("Builtin.hopToActor must be in an async function");
   }
-  auto isolation = getActorIsolationOfContext(FunctionDC);
+  auto isolation =
+      getActorIsolationOfContext(FunctionDC, [](AbstractClosureExpr *CE) {
+        return CE->getActorIsolation();
+      });
   if (isolation != ActorIsolation::Independent
       && isolation != ActorIsolation::Unspecified) {
     // TODO: Explicit hop with no hop-back should only be allowed in independent

lib/Sema/TypeCheckConcurrency.cpp

@@ -1477,10 +1477,14 @@ static bool checkedByFlowIsolation(DeclContext const *refCxt,
 }
 
 /// Get the actor isolation of the innermost relevant context.
-static ActorIsolation getInnermostIsolatedContext(const DeclContext *dc) {
+static ActorIsolation getInnermostIsolatedContext(
+    const DeclContext *dc,
+    llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+        getClosureActorIsolation) {
   // Retrieve the actor isolation of the context.
   auto mutableDC = const_cast<DeclContext *>(dc);
-  switch (auto isolation = getActorIsolationOfContext(mutableDC)) {
+  switch (auto isolation =
+              getActorIsolationOfContext(mutableDC, getClosureActorIsolation)) {
   case ActorIsolation::ActorInstance:
   case ActorIsolation::Independent:
   case ActorIsolation::Unspecified:
@@ -1564,6 +1568,9 @@ namespace {
     SmallVector<ApplyExpr*, 4> applyStack;
     SmallVector<std::pair<OpaqueValueExpr *, Expr *>, 4> opaqueValues;
     SmallVector<const PatternBindingDecl *, 2> patternBindingStack;
+    llvm::function_ref<Type(Expr *)> getType;
+    llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+        getClosureActorIsolation;
 
     /// Keeps track of the capture context of variables that have been
     /// explicitly captured in closures.
@@ -1762,7 +1769,13 @@ namespace {
     }
 
   public:
-    ActorIsolationChecker(const DeclContext *dc) : ctx(dc->getASTContext()) {
+    ActorIsolationChecker(
+        const DeclContext *dc,
+        llvm::function_ref<Type(Expr *)> getType = __Expr_getType,
+        llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+            getClosureActorIsolation = __AbstractClosureExpr_getActorIsolation)
+        : ctx(dc->getASTContext()), getType(getType),
+          getClosureActorIsolation(getClosureActorIsolation) {
       contextStack.push_back(dc);
     }
 
@@ -2062,7 +2075,7 @@ namespace {
           break;
 
         if (auto closure = dyn_cast<AbstractClosureExpr>(dc)) {
-          auto isolation = closure->getActorIsolation();
+          auto isolation = getClosureActorIsolation(closure);
           switch (isolation) {
           case ClosureActorIsolation::Independent:
             if (isSendableClosure(closure, /*forActorIsolation=*/true)) {
@@ -2117,7 +2130,8 @@ namespace {
         // Check isolation of the context itself. We do this separately
         // from the closure check because closures capture specific variables
         // while general isolation is declaration-based.
-        switch (auto isolation = getActorIsolationOfContext(dc)) {
+        switch (auto isolation =
+                    getActorIsolationOfContext(dc, getClosureActorIsolation)) {
         case ActorIsolation::Independent:
         case ActorIsolation::Unspecified:
           // Local functions can capture an isolated parameter.
@@ -2491,7 +2505,7 @@ namespace {
 
     /// Check actor isolation for a particular application.
     bool checkApply(ApplyExpr *apply) {
-      auto fnExprType = apply->getFn()->getType();
+      auto fnExprType = getType(apply->getFn());
       if (!fnExprType)
         return false;
 
@@ -2506,7 +2520,8 @@ namespace {
           return *contextIsolation;
 
         auto declContext = const_cast<DeclContext *>(getDeclContext());
-        contextIsolation = getInnermostIsolatedContext(declContext);
+        contextIsolation =
+            getInnermostIsolatedContext(declContext, getClosureActorIsolation);
         return *contextIsolation;
       };
 
@@ -2542,9 +2557,9 @@ namespace {
         // FIXME: The modeling of unsatisfiedIsolation is not great here.
         // We'd be better off using something more like closure isolation
         // that can talk about specific parameters.
-        auto nominal = arg->getType()->getAnyNominal();
+        auto nominal = getType(arg)->getAnyNominal();
         if (!nominal) {
-          nominal = arg->getType()->getASTContext().getProtocol(
+          nominal = getType(arg)->getASTContext().getProtocol(
               KnownProtocolKind::Actor);
         }
 
@@ -2772,7 +2787,7 @@ namespace {
         // where k is a captured dictionary key.
         if (auto *args = component.getSubscriptArgs()) {
           for (auto arg : *args) {
-            auto type = arg.getExpr()->getType();
+            auto type = getType(arg.getExpr());
             if (type &&
                 shouldDiagnoseExistingDataRaces(getDeclContext()) &&
                 diagnoseNonSendableTypes(
@@ -2805,8 +2820,8 @@ namespace {
       if (base)
         isolatedActor.emplace(getIsolatedActor(base));
       auto result = ActorReferenceResult::forReference(
-          declRef, loc, getDeclContext(),
-          kindOfUsage(decl, context), isolatedActor);
+          declRef, loc, getDeclContext(), kindOfUsage(decl, context),
+          isolatedActor, None, None, getClosureActorIsolation);
       switch (result) {
       case ActorReferenceResult::SameConcurrencyDomain:
         if (diagnoseReferenceToUnsafeGlobal(decl, loc))
@@ -2890,7 +2905,8 @@ namespace {
           refKind = isolatedActor->kind;
           refGlobalActor = isolatedActor->globalActor;
         } else {
-          auto contextIsolation = getInnermostIsolatedContext(getDeclContext());
+          auto contextIsolation = getInnermostIsolatedContext(
+              getDeclContext(), getClosureActorIsolation);
           switch (contextIsolation) {
           case ActorIsolation::ActorInstance:
             refKind = ReferencedActor::Isolated;
@@ -2939,7 +2955,7 @@ namespace {
     // Attempt to resolve the global actor type of a closure.
     Type resolveGlobalActorType(ClosureExpr *closure) {
       // Check whether the closure's type has a global actor already.
-      if (Type closureType = closure->getType()) {
+      if (Type closureType = getType(closure)) {
         if (auto closureFnType = closureType->getAs<FunctionType>()) {
           if (Type globalActor = closureFnType->getGlobalActor())
             return globalActor;
@@ -2981,7 +2997,8 @@ namespace {
         return ClosureActorIsolation::forIndependent(preconcurrency);
 
       // A non-Sendable closure gets its isolation from its context.
-      auto parentIsolation = getActorIsolationOfContext(closure->getParent());
+      auto parentIsolation = getActorIsolationOfContext(
+          closure->getParent(), getClosureActorIsolation);
       preconcurrency |= parentIsolation.preconcurrency();
 
       // We must have parent isolation determined to get here.
@@ -3019,10 +3036,10 @@ bool ActorIsolationChecker::mayExecuteConcurrentlyWith(
   // If both contexts are isolated to the same actor, then they will not
   // execute concurrently.
   auto useIsolation = getActorIsolationOfContext(
-      const_cast<DeclContext *>(useContext));
+      const_cast<DeclContext *>(useContext), getClosureActorIsolation);
   if (useIsolation.isActorIsolated()) {
     auto defIsolation = getActorIsolationOfContext(
-        const_cast<DeclContext *>(defContext));
+        const_cast<DeclContext *>(defContext), getClosureActorIsolation);
     if (useIsolation == defIsolation)
       return false;
   }
@@ -3096,9 +3113,12 @@ void swift::checkPropertyWrapperActorIsolation(
   expr->walk(checker);
 }
 
-ClosureActorIsolation
-swift::determineClosureActorIsolation(AbstractClosureExpr *closure) {
-  ActorIsolationChecker checker(closure->getParent());
+ClosureActorIsolation swift::determineClosureActorIsolation(
+    AbstractClosureExpr *closure, llvm::function_ref<Type(Expr *)> getType,
+    llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+        getClosureActorIsolation) {
+  ActorIsolationChecker checker(closure->getParent(), getType,
+                                getClosureActorIsolation);
   return checker.determineClosureIsolation(closure);
 }
 
@@ -5120,10 +5140,11 @@ static bool equivalentIsolationContexts(
 
 ActorReferenceResult ActorReferenceResult::forReference(
     ConcreteDeclRef declRef, SourceLoc declRefLoc, const DeclContext *fromDC,
-    Optional<VarRefUseEnv> useKind,
-    Optional<ReferencedActor> actorInstance,
+    Optional<VarRefUseEnv> useKind, Optional<ReferencedActor> actorInstance,
     Optional<ActorIsolation> knownDeclIsolation,
-    Optional<ActorIsolation> knownContextIsolation) {
+    Optional<ActorIsolation> knownContextIsolation,
+    llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+        getClosureActorIsolation) {
   // If not provided, compute the isolation of the declaration, adjusted
   // for references.
   ActorIsolation declIsolation = ActorIsolation::forUnspecified();
@@ -5145,7 +5166,8 @@ ActorReferenceResult ActorReferenceResult::forReference(
   if (knownContextIsolation) {
     contextIsolation = *knownContextIsolation;
   } else {
-    contextIsolation = getInnermostIsolatedContext(fromDC);
+    contextIsolation =
+        getInnermostIsolatedContext(fromDC, getClosureActorIsolation);
   }
 
   // When the declaration is not actor-isolated, it can always be accessed