Merge pull request #59944 from ahoppen/pr/no-leaveclosurebodiesunchecked

[CodeCompletion] Eliminate LeaveClosureBodiesUnchecked for solver-based code completion

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/ArgumentCompletion.cpp

@@ -88,6 +88,15 @@ void ArgumentTypeCheckCompletionCallback::sawSolutionImpl(const Solution &S) {
   while (ParentCall && ParentCall->getArgs() == nullptr) {
     ParentCall = CS.getParentExpr(ParentCall);
   }
+  if (auto TV = S.getType(CompletionExpr)->getAs<TypeVariableType>()) {
+    auto Locator = TV->getImpl().getLocator();
+    if (Locator->isLastElement<LocatorPathElt::PatternMatch>()) {
+      // The code completion token is inside a pattern, which got rewritten from
+      // a call by ResolvePattern. Thus, we aren't actually inside a call.
+      // Rest 'ParentCall' to nullptr to reflect that.
+      ParentCall = nullptr;
+    }
+  }
 
   if (!ParentCall || ParentCall == CompletionExpr) {
     // We might not have a call that contains the code completion expression if

lib/IDE/PostfixCompletion.cpp

@@ -44,6 +44,32 @@ 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;
+      }
+    }
+    if (!S.hasType(E)) {
+      return Type();
+    }
+    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();
@@ -60,29 +86,45 @@ void PostfixCompletionCallback::sawSolutionImpl(
   auto *Locator = CS.getConstraintLocator(SemanticExpr);
   Type ExpectedTy = getTypeForCompletion(S, CompletionExpr);
   Expr *ParentExpr = CS.getParentExpr(CompletionExpr);
-  if (!ParentExpr)
+  if (!ParentExpr && !ExpectedTy)
     ExpectedTy = CS.getContextualType(CompletionExpr, /*forConstraint=*/false);
 
   auto *CalleeLocator = S.getCalleeLocator(Locator);
   ValueDecl *ReferencedDecl = nullptr;
   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()});
   if (Ret.second) {
     bool ISDMT = S.isStaticallyDerivedMetatype(ParsedExpr);
     bool ImplicitReturn = isImplicitSingleExpressionReturn(CS, CompletionExpr);
-    bool DisallowVoid = ExpectedTy
-                            ? !ExpectedTy->isVoid()
-                            : !ParentExpr && CS.getContextualTypePurpose(
-                                                 CompletionExpr) != CTP_Unused;
+    bool DisallowVoid = false;
+    DisallowVoid |= ExpectedTy && !ExpectedTy->isVoid();
+    DisallowVoid |= !ParentExpr &&
+                    CS.getContextualTypePurpose(CompletionExpr) != CTP_Unused;
+    for (auto SAT : S.solutionApplicationTargets) {
+      if (DisallowVoid) {
+        // DisallowVoid is already set. No need to iterate further.
+        break;
+      }
+      if (SAT.second.getAsExpr() == CompletionExpr) {
+        DisallowVoid |= SAT.second.getExprContextualTypePurpose() != CTP_Unused;
+      }
+    }
 
     Results.push_back({BaseTy, ReferencedDecl,
                        /*ExpectedTypes=*/{}, DisallowVoid, ISDMT,
-                       ImplicitReturn, IsAsync});
+                       ImplicitReturn, IsAsync, ClosureActorIsolations});
     if (ExpectedTy) {
       Results.back().ExpectedTypes.push_back(ExpectedTy);
     }
@@ -122,6 +164,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,

lib/Parse/ParseExpr.cpp

@@ -1175,6 +1175,7 @@ Parser::parseExprPostfixSuffix(ParserResult<Expr> Result, bool isExprBasic,
       //   [.foo(), <HERE> .bar()]
       // '.bar()' is probably not a part of the inserting element. Moreover,
       // having suffixes doesn't help type inference in any way.
+
       return Result;
     }
 

lib/Parse/ParsePattern.cpp

@@ -1303,6 +1303,19 @@ ParserResult<Pattern> Parser::parseMatchingPattern(bool isExprBasic) {
   if (subExpr.isNull())
     return status;
 
+  if (isa<CodeCompletionExpr>(subExpr.get()) && Tok.isFollowingLParen()) {
+    // We are in the case like the following of parsing a pattern with the code
+    // completion token as base and associated value matches:
+    // #^COMPLETE^#(let a)
+    // We will have not consumed the `(let a)` in `parseExprPostfixSuffix`
+    // because usually suffixes don't influence the code completion's type and
+    // the suffix might be unrelated. But the trailing `(let a)` that is left
+    // prevents us from forming a valid pattern.
+    // Consume and discard the `(let a)`, which just leaves us with the base
+    // of the pattern.
+    (void)parseExprCallSuffix(subExpr, isExprBasic);
+  }
+
   // The most common case here is to parse something that was a lexically
   // obvious pattern, which will come back wrapped in an immediate
   // UnresolvedPatternExpr.  Transform this now to simplify later code.

lib/Sema/CSGen.cpp

@@ -2393,6 +2393,12 @@ namespace {
       // function, to set the type of the pattern.
       auto setType = [&](Type type) {
         CS.setType(pattern, type);
+        if (auto PE = dyn_cast<ExprPattern>(pattern)) {
+          // Set the type of the pattern's sub-expression as well, so code
+          // completion can retrieve the expression's type in case it is a code
+          // completion token.
+          CS.setType(PE->getSubExpr(), type);
+        }
         return type;
       };
 

lib/Sema/CSSyntacticElement.cpp

@@ -1030,6 +1030,23 @@ class SyntacticElementConstraintGenerator
 
     SmallVector<ElementInfo, 4> elements;
     for (auto element : braceStmt->getElements()) {
+      if (cs.isForCodeCompletion() &&
+          !cs.containsIDEInspectionTarget(element)) {
+        // Statements and expressions can't influence the expresion that
+        // contains the code completion token. To improve performance, skip
+        // type checking them entirely.
+        if (element.is<Expr *>() && !element.isExpr(ExprKind::TypeJoin)) {
+          // Type join expressions are not really pure expressions, they kind of
+          // declare new type variables and are important to a result builder's
+          // structure. Don't skip them.
+          continue;
+        } else if (element.is<Stmt *>() && !element.isStmt(StmtKind::Guard)) {
+          // Guard statements might define variables that are used in the code
+          // completion expression. Don't skip them.
+          continue;
+        }
+      }
+
       if (auto *decl = element.dyn_cast<Decl *>()) {
         if (auto *PDB = dyn_cast<PatternBindingDecl>(decl)) {
           visitPatternBinding(PDB, elements);

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -307,7 +307,9 @@ getTypeOfExpressionWithoutApplying(Expr *&expr, DeclContext *dc,
 
   ConstraintSystemOptions options;
   options |= ConstraintSystemFlags::SuppressDiagnostics;
-  options |= ConstraintSystemFlags::LeaveClosureBodyUnchecked;
+  if (!Context.CompletionCallback) {
+    options |= ConstraintSystemFlags::LeaveClosureBodyUnchecked;
+  }
 
   // Construct a constraint system from this expression.
   ConstraintSystem cs(dc, options);
@@ -400,7 +402,6 @@ getTypeOfCompletionOperatorImpl(DeclContext *DC, Expr *expr,
 
   ConstraintSystemOptions options;
   options |= ConstraintSystemFlags::SuppressDiagnostics;
-  options |= ConstraintSystemFlags::LeaveClosureBodyUnchecked;
 
   // Construct a constraint system from this expression.
   ConstraintSystem CS(DC, options);
@@ -612,8 +613,9 @@ bool TypeChecker::typeCheckForCodeCompletion(
     options |= ConstraintSystemFlags::AllowFixes;
     options |= ConstraintSystemFlags::SuppressDiagnostics;
     options |= ConstraintSystemFlags::ForCodeCompletion;
-    options |= ConstraintSystemFlags::LeaveClosureBodyUnchecked;
-
+    if (!Context.CompletionCallback) {
+      options |= ConstraintSystemFlags::LeaveClosureBodyUnchecked;
+    }
 
     ConstraintSystem cs(DC, options);
 

lib/Sema/TypeCheckStmt.cpp

@@ -2393,6 +2393,8 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
     }
   }
 
+  bool LeaveBodyUnchecked = !ctx.CompletionCallback;
+
   // The enclosing closure might be a single expression closure or a function
   // builder closure. In such cases, the body elements are type checked with
   // the closure itself. So we need to try type checking the enclosing closure
@@ -2416,13 +2418,17 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
       auto ActorIsolation = determineClosureActorIsolation(
           CE, __Expr_getType, __AbstractClosureExpr_getActorIsolation);
       CE->setActorIsolation(ActorIsolation);
+      if (!LeaveBodyUnchecked) {
+        // Type checking the parent closure also type checked this node.
+        // Nothing to do anymore.
+        return false;
+      }
       if (CE->getBodyState() != ClosureExpr::BodyState::ReadyForTypeChecking)
         return false;
     }
   }
 
-  TypeChecker::typeCheckASTNode(finder.getRef(), DC,
-                                /*LeaveBodyUnchecked=*/true);
+  TypeChecker::typeCheckASTNode(finder.getRef(), DC, LeaveBodyUnchecked);
   return false;
 }
 

test/IDE/complete_in_closures.swift

@@ -104,7 +104,7 @@ struct NestedStructWithClosureMember1 {
 struct StructWithClosureMemberAndLocal {
   var c = {
     var x = 0
-    #^DELAYED_10?check=WITH_GLOBAL_DECLS_AND_LOCAL1;xfail=https://github.com/apple/swift/issues/58273^#
+    #^DELAYED_10?check=WITH_GLOBAL_DECLS_AND_LOCAL1^#
   }
 }
 
@@ -416,3 +416,83 @@ func testClosureInPatternBindingInit() {
   // CLOSURE_IN_PATTERN_BINDING: End completion
 
 }
+
+func testSwitchInClosure() {
+  func executeClosure(closure: () -> Void) {}
+
+  struct Boredom {
+    static func doNothing() {}
+  }
+
+  enum MyEnum {
+    case first
+    case second
+  }
+
+  let item: MyEnum? = nil
+  executeClosure(closure: {
+    switch item {
+    case .#^SWITCH_IN_CLOSURE^#first:
+      break
+    case .second:
+      Boredom.doNothing()
+    }
+  })
+
+// SWITCH_IN_CLOSURE: Begin completions
+// SWITCH_IN_CLOSURE-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: first[#MyEnum#];
+// SWITCH_IN_CLOSURE-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: second[#MyEnum#];
+// SWITCH_IN_CLOSURE: End completions
+}
+
+func testSwitchWithAssociatedValueInClosure() {
+  func executeClosure(closure: () -> Void) {}
+
+  struct Boredom {
+    static func doNothing() {}
+  }
+
+  enum MyEnum {
+    case first(String)
+  }
+
+  let item: MyEnum? = nil
+  executeClosure(closure: {
+    switch item {
+    case .#^SWITCH_WITH_ASSOC_IN_CLOSURE^#first(_):
+      break
+    }
+  })
+
+// SWITCH_WITH_ASSOC_IN_CLOSURE: Begin completions
+// SWITCH_WITH_ASSOC_IN_CLOSURE-DAG: Decl[EnumElement]/CurrNominal/Flair[ExprSpecific]/TypeRelation[Convertible]: first({#String#})[#MyEnum#];
+// SWITCH_WITH_ASSOC_IN_CLOSURE: End completions
+}
+
+func testCompleteInMatchOfAssociatedValueInSwitchCase() {
+  func testSwitchWithAssociatedValueInClosure() {
+  func executeClosure(closure: () -> Void) {}
+
+  struct Boredom {
+    static func doNothing() {}
+  }
+
+  enum MyEnum {
+    case first(Bool, String)
+  }
+
+  let item: MyEnum? = nil
+  let str = "hi"
+  executeClosure(closure: {
+    switch item {
+    case .first(true, #^IN_ASSOC_OF_CASE_IN_CLOSURE^#str):
+      break
+    }
+  })
+
+// IN_ASSOC_OF_CASE_IN_CLOSURE: Begin completions
+// IN_ASSOC_OF_CASE_IN_CLOSURE-DAG: Decl[LocalVar]/Local:               str[#String#]; name=str
+// IN_ASSOC_OF_CASE_IN_CLOSURE: End completions
+}
+
+}