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

When we get rid of `LeaveClosureBodiesUnchecked` we no longer save closure types to the AST and thus also don’t save their actor isolation to the AST. Hence, we need to extract types and actor isolations of parent closures from the constraint system solution instead of the AST. This prepares `ActorIsolationChecker` to take custom functions to determine the type of an expression or the actor isolation of a closure.

Author: ahoppen

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/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,7 +9229,7 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
     return getActorIsolation(var);
 
   if (auto *closure = dyn_cast<AbstractClosureExpr>(dcToUse)) {
-    return closure->getActorIsolation().getActorIsolation();
+    return getClosureActorIsolation(closure).getActorIsolation();
   }
 
   if (auto *tld = dyn_cast<TopLevelCodeDecl>(dcToUse)) {

lib/AST/Expr.cpp

@@ -1903,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

lib/Sema/TypeCheckConcurrency.h

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ConcreteDeclRef.h"
 #include "swift/AST/DiagnosticEngine.h"
+#include "swift/AST/Expr.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/Type.h"
 #include <cassert>
@@ -60,16 +61,6 @@ void checkInitializerActorIsolation(Initializer *init, Expr *expr);
 void checkEnumElementActorIsolation(EnumElementDecl *element, Expr *expr);
 void checkPropertyWrapperActorIsolation(VarDecl *wrappedVar, Expr *expr);
 
-/// Determine the isolation of a particular closure.
-///
-/// This forwards to \c ActorIsolationChecker::determineClosureActorIsolation
-/// and thus assumes that enclosing closures have already had their isolation
-/// checked.
-///
-/// This does not set the closure's actor isolation
-ClosureActorIsolation
-determineClosureActorIsolation(AbstractClosureExpr *closure);
-
 /// States the reason for checking the Sendability of a given declaration.
 enum class SendableCheckReason {
   /// A reference to an actor from outside that actor.
@@ -244,13 +235,13 @@ struct ActorReferenceResult {
   /// provided when referencing the declaration. This can be either the base
   /// of a member access or a parameter passed to a function.
   static ActorReferenceResult forReference(
-      ConcreteDeclRef declRef,
-      SourceLoc declRefLoc,
-      const DeclContext *fromDC,
+      ConcreteDeclRef declRef, SourceLoc declRefLoc, const DeclContext *fromDC,
       Optional<VarRefUseEnv> useKind = None,
       Optional<ReferencedActor> actorInstance = None,
       Optional<ActorIsolation> knownDeclIsolation = None,
-      Optional<ActorIsolation> knownContextIsolation = None);
+      Optional<ActorIsolation> knownContextIsolation = None,
+      llvm::function_ref<ClosureActorIsolation(AbstractClosureExpr *)>
+          getClosureActorIsolation = __AbstractClosureExpr_getActorIsolation);
 
   operator Kind() const { return kind; }
 };

lib/Sema/TypeCheckStmt.cpp

@@ -1988,7 +1988,8 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
       // isolation checked before nested ones are being checked by the way
       // TypeCheckASTNodeAtLocRequest is called multiple times, as described
       // above.
-      auto ActorIsolation = determineClosureActorIsolation(CE);
+      auto ActorIsolation = determineClosureActorIsolation(
+          CE, __Expr_getType, __AbstractClosureExpr_getActorIsolation);
       CE->setActorIsolation(ActorIsolation);
       if (CE->getBodyState() != ClosureExpr::BodyState::ReadyForTypeChecking)
         return false;