Merge pull request #41883 from ahoppen/pr/solver-based-completion-improvements

[CodeCompletion] Small improvements to solver-based code completion

Author: ahoppen

include/swift/Sema/ConstraintSystem.h

@@ -2415,7 +2415,7 @@ class ConstraintSystem {
 
   /// Arguments after the code completion token that were thus ignored (i.e.
   /// assigned fresh type variables) for type checking.
-  llvm::SetVector<ConstraintLocator *> IgnoredArguments;
+  llvm::SetVector<Expr *> IgnoredArguments;
 
   /// Maps node types used within all portions of the constraint
   /// system, instead of directly using the types on the
@@ -3187,18 +3187,25 @@ class ConstraintSystem {
   bool containsCodeCompletionLoc(ASTNode node) const;
   bool containsCodeCompletionLoc(const ArgumentList *args) const;
 
-  /// Marks the argument with the \p ArgLoc locator as being ignored because it
-  /// occurs after the code completion token. This assumes that the argument is
-  /// not type checked (by assigning it a fresh type variable) and prevents
-  /// fixes from being generated for this argument.
-  void markArgumentIgnoredForCodeCompletion(ConstraintLocator *ArgLoc) {
-    IgnoredArguments.insert(ArgLoc);
+  /// Marks the argument \p Arg as being ignored because it occurs after the
+  /// code completion token. This assumes that the argument is not type checked
+  /// (by assigning it a fresh type variable) and prevents fixes from being
+  /// generated for this argument.
+  void markArgumentIgnoredForCodeCompletion(Expr *Arg) {
+    IgnoredArguments.insert(Arg);
   }
 
-  /// Whether the argument with the \p ArgLoc locator occurs after the code
-  /// completion tokena and thus should be ignored and not generate any fixes.
-  bool isArgumentIgnoredForCodeCompletion(ConstraintLocator *ArgLoc) {
-    return IgnoredArguments.count(ArgLoc) > 0;
+  /// Whether the argument \p Arg occurs after the code completion token and
+  /// thus should be ignored and not generate any fixes.
+  bool isArgumentIgnoredForCodeCompletion(Expr *Arg) const {
+    return IgnoredArguments.count(Arg) > 0;
+  }
+
+  /// Whether the constraint system has ignored any arguments for code
+  /// completion, i.e. whether there is an expression for which
+  /// \c isArgumentIgnoredForCodeCompletion returns \c true.
+  bool hasArgumentsIgnoredForCodeCompletion() const {
+    return !IgnoredArguments.empty();
   }
 
   void setClosureType(const ClosureExpr *closure, FunctionType *type) {

lib/Sema/BuilderTransform.cpp

@@ -1744,6 +1744,7 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
     for (const auto &solution : solutions) {
       cs.getASTContext().CompletionCallback->sawSolution(solution);
     }
+    return nullptr;
   }
 
   if (solvingFailed || solutions.size() != 1) {

lib/Sema/CSBindings.cpp

@@ -2024,9 +2024,13 @@ bool TypeVariableBinding::attempt(ConstraintSystem &cs) const {
         return false;
 
       // Don't penailze solutions that have holes for ignored arguments.
-      if (cs.isArgumentIgnoredForCodeCompletion(
-              TypeVar->getImpl().getLocator())) {
-        return false;
+      if (cs.hasArgumentsIgnoredForCodeCompletion()) {
+        // Avoid simplifying the locator if the constriant system didn't ignore
+        // any arguments.
+        auto argExpr = simplifyLocatorToAnchor(TypeVar->getImpl().getLocator());
+        if (cs.isArgumentIgnoredForCodeCompletion(argExpr.dyn_cast<Expr *>())) {
+          return false;
+        }
       }
     }
     // Reflect in the score that this type variable couldn't be

lib/Sema/CSGen.cpp

@@ -788,8 +788,7 @@ namespace {
       CS(cs) {}
 
     std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
-      if (CS.isArgumentIgnoredForCodeCompletion(
-              CS.getConstraintLocator(expr))) {
+      if (CS.isArgumentIgnoredForCodeCompletion(expr)) {
         return {false, expr};
       }
 
@@ -3647,8 +3646,7 @@ namespace {
     std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
       auto &CS = CG.getConstraintSystem();
 
-      if (CS.isArgumentIgnoredForCodeCompletion(
-              CS.getConstraintLocator(expr))) {
+      if (CS.isArgumentIgnoredForCodeCompletion(expr)) {
         CG.setTypeForArgumentIgnoredForCompletion(expr);
         return {false, expr};
       }
@@ -3723,8 +3721,7 @@ namespace {
         SmallVector<Expr *, 2> ignoredArgs;
         getArgumentsAfterCodeCompletionToken(expr, CS, ignoredArgs);
         for (auto ignoredArg : ignoredArgs) {
-          CS.markArgumentIgnoredForCodeCompletion(
-              CS.getConstraintLocator(ignoredArg));
+          CS.markArgumentIgnoredForCodeCompletion(ignoredArg);
         }
       }
 

lib/Sema/CSSimplify.cpp

@@ -12470,9 +12470,14 @@ bool ConstraintSystem::recordFix(ConstraintFix *fix, unsigned impact) {
     log << ")\n";
   }
 
-  if (isArgumentIgnoredForCodeCompletion(fix->getLocator())) {
-    // The argument was ignored. Don't record any fixes for it.
-    return false;
+  if (hasArgumentsIgnoredForCodeCompletion()) {
+    // Avoid simplifying the locator if the constraint system didn't ignore any
+    // arguments.
+    auto argExpr = simplifyLocatorToAnchor(fix->getLocator());
+    if (isArgumentIgnoredForCodeCompletion(getAsExpr<Expr>(argExpr))) {
+      // The argument was ignored. Don't record any fixes for it.
+      return false;
+    }
   }
 
   // Record the fix.