[CodeCompletion] Record fixes while solving result builders for code completion

We record fixes while solving normal expressions for code completion and we should do the same when solving result builders if we are reporting the solutions to completion callbacks.

Author: ahoppen

include/swift/Sema/ConstraintSystem.h

@@ -5296,18 +5296,21 @@ class ConstraintSystem {
                                  = FreeTypeVariableBinding::Disallow,
                                  bool allowFixes = false);
 
-  /// Construct and solve a system of constraints based on the given expression
-  /// and its contextual information.
+  /// Assuming that constraints have already been generated, solve the
+  /// constraint system for code completion, writing all solutions to
+  /// \p solutions.
   ///
   /// This method is designed to be used for code completion which means that
   /// it doesn't mutate given expression, even if there is a single valid
   /// solution, and constraint solver is allowed to produce partially correct
   /// solutions. Such solutions can have any number of holes in them.
   ///
-  /// \param target The expression involved in code completion.
-  ///
   /// \param solutions The solutions produced for the given target without
   /// filtering.
+  void solveForCodeCompletion(SmallVectorImpl<Solution> &solutions);
+
+  /// Generate constraints for \p target and solve the resulting constraint
+  /// system for code completion (see overload above).
   ///
   /// \returns `false` if this call fails (e.g. pre-check or constraint
   /// generation fails), `true` otherwise.

lib/Sema/BuilderTransform.cpp

@@ -1714,10 +1714,10 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
   }
 
   // Solve the constraint system.
-  SmallVector<Solution, 4> solutions;
-  bool solvingFailed = cs.solve(solutions);
-
   if (cs.getASTContext().CompletionCallback) {
+    SmallVector<Solution, 4> solutions;
+    cs.solveForCodeCompletion(solutions);
+
     CompletionContextFinder analyzer(func, func->getDeclContext());
     filterSolutionsForCodeCompletion(solutions, analyzer);
     for (const auto &solution : solutions) {
@@ -1726,6 +1726,9 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
     return nullptr;
   }
 
+  SmallVector<Solution, 4> solutions;
+  bool solvingFailed = cs.solve(solutions);
+
   if (solvingFailed || solutions.size() != 1) {
     // Try to fix the system or provide a decent diagnostic.
     auto salvagedResult = cs.salvage();

lib/Sema/CSSolver.cpp

@@ -1516,26 +1516,8 @@ void ConstraintSystem::solveImpl(SmallVectorImpl<Solution> &solutions) {
   }
 }
 
-bool ConstraintSystem::solveForCodeCompletion(
-    SolutionApplicationTarget &target, SmallVectorImpl<Solution> &solutions) {
-  auto *expr = target.getAsExpr();
-  // Tell the constraint system what the contextual type is.
-  setContextualType(expr, target.getExprContextualTypeLoc(),
-                    target.getExprContextualTypePurpose());
-
-  // Set up the expression type checker timer.
-  Timer.emplace(expr, *this);
-
-  shrink(expr);
-
-  if (isDebugMode()) {
-    auto &log = llvm::errs();
-    log << "--- Code Completion ---\n";
-  }
-
-  if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
-    return false;
-
+void ConstraintSystem::solveForCodeCompletion(
+    SmallVectorImpl<Solution> &solutions) {
   {
     SolverState state(*this, FreeTypeVariableBinding::Disallow);
 
@@ -1556,6 +1538,30 @@ bool ConstraintSystem::solveForCodeCompletion(
     }
   }
 
+  return;
+}
+
+bool ConstraintSystem::solveForCodeCompletion(
+    SolutionApplicationTarget &target, SmallVectorImpl<Solution> &solutions) {
+  auto *expr = target.getAsExpr();
+  // Tell the constraint system what the contextual type is.
+  setContextualType(expr, target.getExprContextualTypeLoc(),
+                    target.getExprContextualTypePurpose());
+
+  // Set up the expression type checker timer.
+  Timer.emplace(expr, *this);
+
+  shrink(expr);
+
+  if (isDebugMode()) {
+    auto &log = llvm::errs();
+    log << "--- Code Completion ---\n";
+  }
+
+  if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
+    return false;
+
+  solveForCodeCompletion(solutions);
   return true;
 }