[CodeCompletion] Use type from solution to determine actor isolation

Author: ahoppen

include/swift/IDE/PostfixCompletion.h

@@ -35,6 +35,11 @@ class PostfixCompletionCallback : public TypeCheckCompletionCallback {
     /// Whether the surrounding context is async and thus calling async
     /// functions is supported.
     bool IsInAsyncContext;
+
+    /// Actor isolations that were determined during constraint solving but that
+    /// haven't been saved to the AST.
+    llvm::DenseMap<AbstractClosureExpr *, ClosureActorIsolation>
+        ClosureActorIsolations;
   };
 
   CodeCompletionExpr *CompletionExpr;

lib/IDE/PostfixCompletion.cpp

@@ -44,6 +44,29 @@ void PostfixCompletionCallback::fallbackTypeCheck(DeclContext *DC) {
                              [&](const Solution &S) { sawSolution(S); });
 }
 
+static ClosureActorIsolation
+getClosureActorIsolation(const Solution &S, AbstractClosureExpr *ACE) {
+  auto getType = [&S](Expr *E) -> Type {
+    // Prefer the contextual type of the closure because it might be 'weaker'
+    // than the type determined for the closure by the constraints system. E.g.
+    // the contextual type might have a global actor attribute but because no
+    // methods from that global actor are called in the closure, the closure has
+    // a non-actor type.
+    auto target = S.solutionApplicationTargets.find(dyn_cast<ClosureExpr>(E));
+    if (target != S.solutionApplicationTargets.end()) {
+      if (auto Ty = target->second.getClosureContextualType()) {
+        return Ty;
+      }
+    }
+    return getTypeForCompletion(S, E);
+  };
+  auto getClosureActorIsolationThunk = [&S](AbstractClosureExpr *ACE) {
+    return getClosureActorIsolation(S, ACE);
+  };
+  return determineClosureActorIsolation(ACE, getType,
+                                        getClosureActorIsolationThunk);
+}
+
 void PostfixCompletionCallback::sawSolutionImpl(
     const constraints::Solution &S) {
   auto &CS = S.getConstraintSystem();
@@ -68,7 +91,14 @@ void PostfixCompletionCallback::sawSolutionImpl(
   if (auto SelectedOverload = S.getOverloadChoiceIfAvailable(CalleeLocator))
     ReferencedDecl = SelectedOverload->choice.getDeclOrNull();
 
+  llvm::DenseMap<AbstractClosureExpr *, ClosureActorIsolation>
+      ClosureActorIsolations;
   bool IsAsync = isContextAsync(S, DC);
+  for (auto SAT : S.solutionApplicationTargets) {
+    if (auto ACE = getAsExpr<AbstractClosureExpr>(SAT.second.getAsASTNode())) {
+      ClosureActorIsolations[ACE] = getClosureActorIsolation(S, ACE);
+    }
+  }
 
   auto Key = std::make_pair(BaseTy, ReferencedDecl);
   auto Ret = BaseToSolutionIdx.insert({Key, Results.size()});
@@ -91,7 +121,7 @@ void PostfixCompletionCallback::sawSolutionImpl(
 
     Results.push_back({BaseTy, ReferencedDecl,
                        /*ExpectedTypes=*/{}, DisallowVoid, ISDMT,
-                       ImplicitReturn, IsAsync});
+                       ImplicitReturn, IsAsync, ClosureActorIsolations});
     if (ExpectedTy) {
       Results.back().ExpectedTypes.push_back(ExpectedTy);
     }
@@ -131,6 +161,7 @@ void PostfixCompletionCallback::deliverResults(
   Lookup.shouldCheckForDuplicates(Results.size() > 1);
   for (auto &Result : Results) {
     Lookup.setCanCurrDeclContextHandleAsync(Result.IsInAsyncContext);
+    Lookup.setClosureActorIsolations(Result.ClosureActorIsolations);
     Lookup.setIsStaticMetatype(Result.BaseIsStaticMetaType);
     Lookup.getPostfixKeywordCompletions(Result.BaseTy, BaseExpr);
     Lookup.setExpectedTypes(Result.ExpectedTypes,