[TypeChecker/CodeCompletion] Allow `{typeCheck, solve}ForCodeCompleti…

…on` to operate on target instead of expression

Using `SolutionApplicationTarget` make it easier to propage
contextual information and avoid mistakes of using incorrect
accessors for constraint generation.

lib/Sema/CSSolver.cpp

@@ -1467,40 +1467,45 @@ void ConstraintSystem::solveImpl(SmallVectorImpl<Solution> &solutions) {
 }
 
 bool ConstraintSystem::solveForCodeCompletion(
-    Expr *expr, DeclContext *DC, Type contextualType, ContextualTypePurpose CTP,
+    SolutionApplicationTarget &target,
     llvm::function_ref<void(const Solution &)> callback) {
-  // First, pre-check the expression, validating any types that occur in the
-  // expression and folding sequence expressions.
-  if (ConstraintSystem::preCheckExpression(
-          expr, DC, /*replaceInvalidRefsWithErrors=*/true))
-    return false;
+  auto *expr = target.getAsExpr();
+  {
+    // First, pre-check the expression, validating any types that occur in the
+    // expression and folding sequence expressions.
+    auto failedPreCheck = ConstraintSystem::preCheckExpression(
+        expr, target.getDeclContext(),
+        /*replaceInvalidRefsWithErrors=*/true);
+
+    target.setExpr(expr);
+
+    if (failedPreCheck)
+      return false;
+  }
 
   ConstraintSystemOptions options;
   options |= ConstraintSystemFlags::AllowFixes;
   options |= ConstraintSystemFlags::SuppressDiagnostics;
 
-  ConstraintSystem cs(DC, options);
+  ConstraintSystem cs(target.getDeclContext(), options);
 
-  if (CTP != ContextualTypePurpose::CTP_Unused)
-    cs.setContextualType(expr, TypeLoc::withoutLoc(contextualType), CTP);
+  // Tell the constraint system what the contextual type is.
+  cs.setContextualType(expr, target.getExprContextualTypeLoc(),
+                       target.getExprContextualTypePurpose());
 
   // Set up the expression type checker timer.
   cs.Timer.emplace(expr, cs);
 
   cs.shrink(expr);
 
-  expr = cs.generateConstraints(expr, DC);
-  if (!expr)
-    return false;
-
   if (cs.isDebugMode()) {
     auto &log = llvm::errs();
     log << "--- Code Completion ---\n";
-    cs.print(log, expr);
-    log << '\n';
-    cs.print(log);
   }
 
+  if (cs.generateConstraints(target, FreeTypeVariableBinding::Disallow))
+    return false;
+
   llvm::SmallVector<Solution, 4> solutions;
 
   {

lib/Sema/ConstraintSystem.h

@@ -5011,23 +5011,15 @@ class ConstraintSystem {
   /// solution, and constraint solver is allowed to produce partially correct
   /// solutions. Such solutions can have any number of holes in them.
   ///
-  /// \param expr The expression involved in code completion.
-  ///
-  /// \param DC The declaration context this expression is found in.
-  ///
-  /// \param contextualType The type \p expr is being converted to.
-  ///
-  /// \param CTP When contextualType is specified, this indicates what
-  /// the conversion is doing.
+  /// \param target The expression involved in code completion.
   ///
   /// \param callback The callback to be used to provide results to
   /// code completion.
   ///
   /// \returns `false` if this call fails (e.g. pre-check or constraint
   /// generation fails), `true` otherwise.
   static bool
-  solveForCodeCompletion(Expr *expr, DeclContext *DC, Type contextualType,
-                         ContextualTypePurpose CTP,
+  solveForCodeCompletion(SolutionApplicationTarget &target,
                          llvm::function_ref<void(const Solution &)> callback);
 
 private:

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -593,10 +593,12 @@ TypeChecker::getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
 }
 
 void TypeChecker::typeCheckForCodeCompletion(
-    Expr *expr, DeclContext *DC, Type contextualType, ContextualTypePurpose CTP,
+    SolutionApplicationTarget &target,
     llvm::function_ref<void(const Solution &)> callback) {
+  auto *DC = target.getDeclContext();
   auto &Context = DC->getASTContext();
 
+  auto *expr = target.getAsExpr();
   FrontendStatsTracer StatsTracer(Context.Stats,
                                   "typecheck-for-code-completion", expr);
   PrettyStackTraceExpr stackTrace(Context, "code-completion", expr);
@@ -627,14 +629,16 @@ void TypeChecker::typeCheckForCodeCompletion(
   assert(completionExpr);
 
   // If it was possible to solve for a while expression, we are done.
-  if (ConstraintSystem::solveForCodeCompletion(expr, DC, contextualType, CTP,
-                                               callback))
+  if (ConstraintSystem::solveForCodeCompletion(target, callback))
     return;
 
   // If initial solve failed, let's fallback to checking only code completion
   // expresion without any context.
-  (void)ConstraintSystem::solveForCodeCompletion(completionExpr, DC, Type(),
-                                                 CTP_Unused, callback);
+  SolutionApplicationTarget completionTarget(completionExpr, DC, CTP_Unused,
+                                             /*contextualType=*/Type(),
+                                             /*isDiscarded=*/true);
+
+  (void)ConstraintSystem::solveForCodeCompletion(completionTarget, callback);
 }
 
 static Optional<Type> getTypeOfCompletionContextExpr(
@@ -743,11 +747,10 @@ bool DotExprLookup::isApplicable(Expr *E) {
 
 void DotExprLookup::fallbackTypeCheck() {
   assert(!gotCallback());
-  TypeChecker::typeCheckForCodeCompletion(CompletionExpr, DC, Type(),
-                                          ContextualTypePurpose::CTP_Unused,
-                                          [&](const Solution &S) {
-    sawSolution(S);
-  });
+  SolutionApplicationTarget completionTarget(CompletionExpr, DC, CTP_Unused,
+                                             Type(), /*isDiscared=*/true);
+  TypeChecker::typeCheckForCodeCompletion(
+      completionTarget, [&](const Solution &S) { sawSolution(S); });
 }
 
 void DotExprLookup::sawSolution(const constraints::Solution &S) {

lib/Sema/TypeCheckConstraints.cpp

@@ -2164,12 +2164,9 @@ TypeChecker::typeCheckExpression(
 
   if (Context.CompletionCallback &&
       Context.CompletionCallback->isApplicable(expr)) {
-    typeCheckForCodeCompletion(expr, dc, target.getExprContextualType(),
-                               target.getExprContextualTypePurpose(),
-                               [&](const constraints::Solution &S) {
+    typeCheckForCodeCompletion(target, [&](const constraints::Solution &S) {
       Context.CompletionCallback->sawSolution(S);
     });
-    target.setExpr(expr);
     return None;
   }
 

lib/Sema/TypeChecker.h

@@ -638,7 +638,7 @@ FunctionType *getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
 /// solution, and constraint solver is allowed to produce partially correct
 /// solutions. Such solutions can have any number of holes in them.
 void typeCheckForCodeCompletion(
-    Expr *expr, DeclContext *DC, Type contextualType, ContextualTypePurpose CTP,
+    constraints::SolutionApplicationTarget &target,
     llvm::function_ref<void(const constraints::Solution &)> callback);
 
 /// Check the key-path expression.