[CodeCompletion][Sema] Allow missing args when solving if the completion location indicates the user may intend to write them later.

func foo(a: Int, b: Int) {}
func foo(a: String) {}

// Int and String should both be valid, despite the missing argument for the
// first overload since the second arg may just have not been written yet.
foo(a: <complete here>

func bar(a: (Int) -> ()) {}
func bar(a: (String, Int) -> ()) {}

// $0 being of type String should be valid, rather than just Int, since $1 may
// just have not been written yet.
bar { $0.<complete here> }

Author: Nathan Hawes

include/swift/Sema/CSFix.h

@@ -1178,18 +1178,21 @@ class AllowClosureParamDestructuring final : public ConstraintFix {
                                                 ConstraintLocator *locator);
 };
 
+struct SynthesizedArg {
+  unsigned paramIdx;
+  AnyFunctionType::Param param;
+};
+
 class AddMissingArguments final
     : public ConstraintFix,
       private llvm::TrailingObjects<
-          AddMissingArguments, std::pair<unsigned, AnyFunctionType::Param>> {
+          AddMissingArguments, SynthesizedArg> {
   friend TrailingObjects;
 
-  using SynthesizedParam = std::pair<unsigned, AnyFunctionType::Param>;
-
   unsigned NumSynthesized;
 
   AddMissingArguments(ConstraintSystem &cs,
-                      ArrayRef<SynthesizedParam> synthesizedArgs,
+                      ArrayRef<SynthesizedArg> synthesizedArgs,
                       ConstraintLocator *locator)
       : ConstraintFix(cs, FixKind::AddMissingArguments, locator),
         NumSynthesized(synthesizedArgs.size()) {
@@ -1200,8 +1203,8 @@ class AddMissingArguments final
 public:
   std::string getName() const override { return "synthesize missing argument(s)"; }
 
-  ArrayRef<SynthesizedParam> getSynthesizedArguments() const {
-    return {getTrailingObjects<SynthesizedParam>(), NumSynthesized};
+  ArrayRef<SynthesizedArg> getSynthesizedArguments() const {
+    return {getTrailingObjects<SynthesizedArg>(), NumSynthesized};
   }
 
   bool diagnose(const Solution &solution, bool asNote = false) const override;
@@ -1211,12 +1214,16 @@ class AddMissingArguments final
   }
 
   static AddMissingArguments *create(ConstraintSystem &cs,
-                                     ArrayRef<SynthesizedParam> synthesizedArgs,
+                                     ArrayRef<SynthesizedArg> synthesizedArgs,
                                      ConstraintLocator *locator);
 
+  static bool classof(const ConstraintFix *fix) {
+    return fix->getKind() == FixKind::AddMissingArguments;
+  }
+
 private:
-  MutableArrayRef<SynthesizedParam> getSynthesizedArgumentsBuf() {
-    return {getTrailingObjects<SynthesizedParam>(), NumSynthesized};
+  MutableArrayRef<SynthesizedArg> getSynthesizedArgumentsBuf() {
+    return {getTrailingObjects<SynthesizedArg>(), NumSynthesized};
   }
 };
 

include/swift/Sema/ConstraintSystem.h

@@ -5393,7 +5393,10 @@ class MatchCallArgumentListener {
   /// indices.
   ///
   /// \param paramIdx The index of the parameter that is missing an argument.
-  virtual Optional<unsigned> missingArgument(unsigned paramIdx);
+  /// \param argInsertIdx The index in the argument list where this argument was
+  /// expected.
+  virtual Optional<unsigned> missingArgument(unsigned paramIdx,
+                                             unsigned argInsertIdx);
 
   /// Indicate that there was no label given when one was expected by parameter.
   ///

lib/Sema/CSDiagnostics.cpp

@@ -4097,9 +4097,9 @@ bool MissingArgumentsFailure::diagnoseAsError() {
 
   interleave(
       SynthesizedArgs,
-      [&](const std::pair<unsigned, AnyFunctionType::Param> &e) {
-        const auto paramIdx = e.first;
-        const auto &arg = e.second;
+      [&](const SynthesizedArg &e) {
+        const auto paramIdx = e.paramIdx;
+        const auto &arg = e.param;
 
         if (arg.hasLabel()) {
           arguments << "'" << arg.getLabel().str() << "'";
@@ -4126,8 +4126,8 @@ bool MissingArgumentsFailure::diagnoseAsError() {
     llvm::raw_svector_ostream fixIt(scratch);
     interleave(
         SynthesizedArgs,
-        [&](const std::pair<unsigned, AnyFunctionType::Param> &arg) {
-          forFixIt(fixIt, arg.second);
+        [&](const SynthesizedArg &arg) {
+          forFixIt(fixIt, arg.param);
         },
         [&] { fixIt << ", "; });
 
@@ -4176,8 +4176,8 @@ bool MissingArgumentsFailure::diagnoseSingleMissingArgument() const {
     return false;
 
   const auto &argument = SynthesizedArgs.front();
-  auto position = argument.first;
-  auto label = argument.second.getLabel();
+  auto position = argument.paramIdx;
+  auto label = argument.param.getLabel();
 
   Expr *fnExpr = nullptr;
   Expr *argExpr = nullptr;
@@ -4192,7 +4192,7 @@ bool MissingArgumentsFailure::diagnoseSingleMissingArgument() const {
   }
 
   // Will the parameter accept a trailing closure?
-  Type paramType = resolveType(argument.second.getPlainType());
+  Type paramType = resolveType(argument.param.getPlainType());
   bool paramAcceptsTrailingClosure = paramType
       ->lookThroughAllOptionalTypes()->is<AnyFunctionType>();
 
@@ -4208,7 +4208,7 @@ bool MissingArgumentsFailure::diagnoseSingleMissingArgument() const {
   else if (position != 0)
     insertText << ", ";
 
-  forFixIt(insertText, argument.second);
+  forFixIt(insertText, argument.param);
 
   if (position == 0 && numArgs > 0 &&
       (!firstTrailingClosure || position < *firstTrailingClosure))
@@ -6036,7 +6036,7 @@ bool ArgumentMismatchFailure::diagnoseMisplacedMissingArgument() const {
   auto anchor = getRawAnchor();
 
   MissingArgumentsFailure failure(
-      solution, {std::make_pair(0, param)},
+      solution, {SynthesizedArg{0, param}},
       getConstraintLocator(anchor, ConstraintLocator::ApplyArgument));
 
   return failure.diagnoseSingleMissingArgument();

lib/Sema/CSDiagnostics.h

@@ -1271,13 +1271,11 @@ class ImplicitInitOnNonConstMetatypeFailure final
 };
 
 class MissingArgumentsFailure final : public FailureDiagnostic {
-  using SynthesizedParam = std::pair<unsigned, AnyFunctionType::Param>;
-
-  SmallVector<SynthesizedParam, 4> SynthesizedArgs;
+  SmallVector<SynthesizedArg, 4> SynthesizedArgs;
 
 public:
   MissingArgumentsFailure(const Solution &solution,
-                          ArrayRef<SynthesizedParam> synthesizedArgs,
+                          ArrayRef<SynthesizedArg> synthesizedArgs,
                           ConstraintLocator *locator)
       : FailureDiagnostic(solution, locator),
         SynthesizedArgs(synthesizedArgs.begin(), synthesizedArgs.end()) {

lib/Sema/CSFix.cpp

@@ -748,9 +748,9 @@ bool AddMissingArguments::diagnose(const Solution &solution,
 
 AddMissingArguments *
 AddMissingArguments::create(ConstraintSystem &cs,
-                            ArrayRef<SynthesizedParam> synthesizedArgs,
+                            ArrayRef<SynthesizedArg> synthesizedArgs,
                             ConstraintLocator *locator) {
-  unsigned size = totalSizeToAlloc<SynthesizedParam>(synthesizedArgs.size());
+  unsigned size = totalSizeToAlloc<SynthesizedArg>(synthesizedArgs.size());
   void *mem = cs.getAllocator().Allocate(size, alignof(AddMissingArguments));
   return new (mem) AddMissingArguments(cs, synthesizedArgs, locator);
 }

lib/Sema/CSGen.cpp

@@ -589,6 +589,13 @@ namespace {
 
       // Determine whether the given declaration is favored.
       auto isFavoredDecl = [&](ValueDecl *value, Type type) -> bool {
+        // We want to consider all options for calls that might contain the code
+        // completion location, as missing arguments after the completion
+        // location are valid (since it might be that they just haven't been
+        // written yet).
+        if (CS.isForCodeCompletion())
+          return false;
+
         if (!type->is<AnyFunctionType>())
           return false;
 

lib/Sema/CSSimplify.cpp

@@ -40,7 +40,8 @@ MatchCallArgumentListener::~MatchCallArgumentListener() { }
 bool MatchCallArgumentListener::extraArgument(unsigned argIdx) { return true; }
 
 Optional<unsigned>
-MatchCallArgumentListener::missingArgument(unsigned paramIdx) {
+MatchCallArgumentListener::missingArgument(unsigned paramIdx,
+                                           unsigned argInsertIdx) {
   return None;
 }
 
@@ -701,11 +702,14 @@ static bool matchCallArgumentsImpl(
   }
 
   // If we have any unfulfilled parameters, check them now.
+  Optional<unsigned> prevArgIdx;
   if (haveUnfulfilledParams) {
     for (auto paramIdx : indices(params)) {
       // If we have a binding for this parameter, we're done.
-      if (!parameterBindings[paramIdx].empty())
+      if (!parameterBindings[paramIdx].empty()) {
+        prevArgIdx = parameterBindings[paramIdx].back();
         continue;
+      }
 
       const auto &param = params[paramIdx];
 
@@ -717,7 +721,8 @@ static bool matchCallArgumentsImpl(
       if (paramInfo.hasDefaultArgument(paramIdx))
         continue;
 
-      if (auto newArgIdx = listener.missingArgument(paramIdx)) {
+      unsigned argInsertIdx = prevArgIdx ? *prevArgIdx + 1 : 0;
+      if (auto newArgIdx = listener.missingArgument(paramIdx, argInsertIdx)) {
         parameterBindings[paramIdx].push_back(*newArgIdx);
         continue;
       }
@@ -980,14 +985,47 @@ constraints::matchCallArguments(
   };
 }
 
+static Optional<unsigned>
+getCompletionArgIndex(ASTNode anchor, SourceManager &SM) {
+  Expr *arg = nullptr;
+  if (auto *CE = getAsExpr<CallExpr>(anchor))
+    arg = CE->getArg();
+  if (auto *SE = getAsExpr<SubscriptExpr>(anchor))
+    arg = SE->getIndex();
+  if (auto *OLE = getAsExpr<ObjectLiteralExpr>(anchor))
+    arg =  OLE->getArg();
+
+  if (!arg)
+    return None;
+
+  auto containsCompletion = [&](Expr *elem) {
+    if (!elem)
+      return false;
+    SourceRange range = elem->getSourceRange();
+    return range.isValid() && SM.rangeContainsCodeCompletionLoc(range);
+  };
+
+  if (auto *TE = dyn_cast<TupleExpr>(arg)) {
+    auto elems = TE->getElements();
+    auto idx = llvm::find_if(elems, containsCompletion);
+    if (idx != elems.end())
+      return std::distance(elems.begin(), idx);
+  } else if (auto *PE = dyn_cast<ParenExpr>(arg)) {
+    if (containsCompletion(PE->getSubExpr()))
+      return 0;
+  }
+  return None;
+}
+
 class ArgumentFailureTracker : public MatchCallArgumentListener {
   ConstraintSystem &CS;
   SmallVectorImpl<AnyFunctionType::Param> &Arguments;
   ArrayRef<AnyFunctionType::Param> Parameters;
   ConstraintLocatorBuilder Locator;
 
-  SmallVector<std::pair<unsigned, AnyFunctionType::Param>, 4> MissingArguments;
+  SmallVector<SynthesizedArg, 4> MissingArguments;
   SmallVector<std::pair<unsigned, AnyFunctionType::Param>, 4> ExtraArguments;
+  Optional<unsigned> CompletionArgIdx;
 
 public:
   ArgumentFailureTracker(ConstraintSystem &cs,
@@ -1006,7 +1044,8 @@ class ArgumentFailureTracker : public MatchCallArgumentListener {
     }
   }
 
-  Optional<unsigned> missingArgument(unsigned paramIdx) override {
+  Optional<unsigned> missingArgument(unsigned paramIdx,
+                                     unsigned argInsertIdx) override {
     if (!CS.shouldAttemptFixes())
       return None;
 
@@ -1023,10 +1062,22 @@ class ArgumentFailureTracker : public MatchCallArgumentListener {
                                       TVO_CanBindToNoEscape | TVO_CanBindToHole);
 
     auto synthesizedArg = param.withType(argType);
-
-    MissingArguments.push_back(std::make_pair(paramIdx, synthesizedArg));
     Arguments.push_back(synthesizedArg);
 
+    if (CS.isForCodeCompletion()) {
+      // When solving for code completion, if any argument contains the
+      // completion location, later arguments shouldn't be considered missing
+      // (causing the solution to have a worse score) as the user just hasn't
+      // written them yet. Early exit to avoid recording them in this case.
+      SourceManager &SM = CS.getASTContext().SourceMgr;
+      if (!CompletionArgIdx)
+        CompletionArgIdx = getCompletionArgIndex(Locator.getAnchor(), SM);
+      if (CompletionArgIdx && *CompletionArgIdx < argInsertIdx)
+        return newArgIdx;
+    }
+
+    MissingArguments.push_back(SynthesizedArg{paramIdx, synthesizedArg});
+
     return newArgIdx;
   }
 
@@ -1190,12 +1241,11 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
       argsWithLabels.pop_back();
       // Let's make sure that labels associated with tuple elements
       // line up with what is expected by argument list.
-      SmallVector<std::pair<unsigned, AnyFunctionType::Param>, 4>
-          synthesizedArgs;
+      SmallVector<SynthesizedArg, 4> synthesizedArgs;
       for (unsigned i = 0, n = argTuple->getNumElements(); i != n; ++i) {
         const auto &elt = argTuple->getElement(i);
         AnyFunctionType::Param argument(elt.getType(), elt.getName());
-        synthesizedArgs.push_back(std::make_pair(i, argument));
+        synthesizedArgs.push_back(SynthesizedArg{i, argument});
         argsWithLabels.push_back(argument);
       }
 
@@ -1771,15 +1821,27 @@ static bool fixMissingArguments(ConstraintSystem &cs, ASTNode anchor,
         cs.createTypeVariable(argLoc, TVO_CanBindToNoEscape)));
   }
 
-  SmallVector<std::pair<unsigned, AnyFunctionType::Param>, 4> synthesizedArgs;
+  SmallVector<SynthesizedArg, 4> synthesizedArgs;
   synthesizedArgs.reserve(numMissing);
   for (unsigned i = args.size() - numMissing, n = args.size(); i != n; ++i) {
-    synthesizedArgs.push_back(std::make_pair(i, args[i]));
+    synthesizedArgs.push_back(SynthesizedArg{i, args[i]});
+  }
+
+  // Treat missing anonymous arguments as valid in closures containing the
+  // code completion location, since they may have just not been written yet.
+  if (cs.isForCodeCompletion()) {
+    if (auto *closure = getAsExpr<ClosureExpr>(anchor)) {
+      SourceManager &SM = closure->getASTContext().SourceMgr;
+      SourceRange range = closure->getSourceRange();
+      if (range.isValid() && SM.rangeContainsCodeCompletionLoc(range) &&
+          (closure->hasAnonymousClosureVars() ||
+           (args.empty() && closure->getInLoc().isInvalid())))
+          return false;
+    }
   }
 
   auto *fix = AddMissingArguments::create(cs, synthesizedArgs,
                                           cs.getConstraintLocator(locator));
-
   if (cs.recordFix(fix))
     return true;
 
@@ -3966,7 +4028,7 @@ bool ConstraintSystem::repairFailures(
     // to diagnose this as a missing argument which can't be ignored.
     if (arg != getTypeVariables().end()) {
       conversionsOrFixes.push_back(AddMissingArguments::create(
-          *this, {std::make_pair(0, AnyFunctionType::Param(*arg))},
+          *this, {SynthesizedArg{0, AnyFunctionType::Param(*arg)}},
           getConstraintLocator(anchor, path)));
       break;
     }