[CodeCompletion] Call sawSolution when type checking a result builder

include/swift/Sema/CompletionContextFinder.h

@@ -13,6 +13,7 @@
 #ifndef SWIFT_SEMA_COMPLETIONCONTEXTFINDER_H
 #define SWIFT_SEMA_COMPLETIONCONTEXTFINDER_H
 
+#include "swift/AST/ASTNode.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Expr.h"
 #include "swift/Sema/CodeCompletionTypeChecking.h"
@@ -49,10 +50,10 @@ class CompletionContextFinder : public ASTWalker {
 
 public:
   /// Finder for completion contexts within the provided initial expression.
-  CompletionContextFinder(Expr *initialExpr, DeclContext *DC)
-      : InitialExpr(initialExpr), InitialDC(DC) {
+  CompletionContextFinder(ASTNode initialNode, DeclContext *DC)
+      : InitialExpr(initialNode.dyn_cast<Expr *>()), InitialDC(DC) {
     assert(DC);
-    initialExpr->walk(*this);
+    initialNode.walk(*this);
   };
 
   /// Finder for completion contexts within the outermost non-closure context of

lib/Sema/BuilderTransform.cpp

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

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -513,11 +513,32 @@ TypeChecker::getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
   }
 }
 
-/// Remove any solutions from the provided vector that both require fixes and
-/// have a score worse than the best.
-static void filterSolutions(SolutionApplicationTarget &target,
-                            SmallVectorImpl<Solution> &solutions,
-                            CodeCompletionExpr *completionExpr) {
+static bool hasTypeForCompletion(Solution &solution,
+                                 CompletionContextFinder &contextAnalyzer) {
+  if (contextAnalyzer.hasCompletionExpr()) {
+    return solution.hasType(contextAnalyzer.getCompletionExpr());
+  } else {
+    assert(contextAnalyzer.hasCompletionKeyPathComponent());
+    return solution.hasType(
+        contextAnalyzer.getKeyPathContainingCompletionComponent(),
+        contextAnalyzer.getKeyPathCompletionComponentIndex());
+  }
+}
+
+void TypeChecker::filterSolutionsForCodeCompletion(
+    SmallVectorImpl<Solution> &solutions,
+    CompletionContextFinder &contextAnalyzer) {
+  // Ignore solutions that didn't end up involving the completion (e.g. due to
+  // a fix to skip over/ignore it).
+  llvm::erase_if(solutions, [&](Solution &S) {
+    if (hasTypeForCompletion(S, contextAnalyzer))
+      return false;
+    // FIXME: Technically this should never happen, but it currently does in
+    // result builder contexts. Re-evaluate if we can assert here when we have
+    // multi-statement closure checking for result builders.
+    return true;
+  });
+
   if (solutions.size() <= 1)
     return;
 
@@ -608,15 +629,6 @@ bool TypeChecker::typeCheckForCodeCompletion(
     if (!cs.solveForCodeCompletion(target, solutions))
       return CompletionResult::Fallback;
 
-    // FIXME: instead of filtering, expose the score and viability to clients.
-    // Remove any solutions that both require fixes and have a score that is
-    // worse than the best.
-    CodeCompletionExpr *completionExpr = nullptr;
-    if (contextAnalyzer.hasCompletionExpr()) {
-      completionExpr = contextAnalyzer.getCompletionExpr();
-    }
-    filterSolutions(target, solutions, completionExpr);
-
     // Similarly, if the type-check didn't produce any solutions, fall back
     // to type-checking a sub-expression in isolation.
     if (solutions.empty())
@@ -626,19 +638,7 @@ bool TypeChecker::typeCheckForCodeCompletion(
     // closure body it could either be type-checked together with the context
     // or not, it's impossible to say without checking.
     if (contextAnalyzer.locatedInMultiStmtClosure()) {
-      auto &solution = solutions.front();
-
-      bool HasTypeForCompletionNode = false;
-      if (completionExpr) {
-        HasTypeForCompletionNode = solution.hasType(completionExpr);
-      } else {
-        assert(contextAnalyzer.hasCompletionKeyPathComponent());
-        HasTypeForCompletionNode = solution.hasType(
-            contextAnalyzer.getKeyPathContainingCompletionComponent(),
-            contextAnalyzer.getKeyPathCompletionComponentIndex());
-      }
-
-      if (!HasTypeForCompletionNode) {
+      if (!hasTypeForCompletion(solutions.front(), contextAnalyzer)) {
         // At this point we know the code completion node wasn't checked with
         // the closure's surrounding context, so can defer to regular
         // type-checking for the current call to typeCheckExpression. If that
@@ -651,6 +651,11 @@ bool TypeChecker::typeCheckForCodeCompletion(
       }
     }
 
+    // FIXME: instead of filtering, expose the score and viability to clients.
+    // Remove solutions that skipped over/ignored the code completion point
+    // or that require fixes and have a score that is worse than the best.
+    filterSolutionsForCodeCompletion(solutions, contextAnalyzer);
+
     llvm::for_each(solutions, callback);
     return CompletionResult::Ok;
   };

lib/Sema/TypeChecker.h

@@ -29,10 +29,11 @@
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/TypeRefinementContext.h"
-#include "swift/Parse/Lexer.h"
 #include "swift/Basic/OptionSet.h"
-#include "swift/Sema/ConstraintSystem.h"
 #include "swift/Config.h"
+#include "swift/Parse/Lexer.h"
+#include "swift/Sema/CompletionContextFinder.h"
+#include "swift/Sema/ConstraintSystem.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include <functional>
@@ -571,6 +572,12 @@ FunctionType *getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
                                           DeclRefKind refKind,
                                           ConcreteDeclRef &refdDecl);
 
+/// Remove any solutions from the provided vector that require more fixes than
+/// the best score or don't contain a type for the code completion token.
+void filterSolutionsForCodeCompletion(
+    SmallVectorImpl<constraints::Solution> &solutions,
+    CompletionContextFinder &contextAnalyzer);
+
 /// Type check the given expression and provide results back to code completion
 /// via specified callback.
 ///