[CodeCompletion][Sema] Migrate CallArgurment position completion to the solver-based implementation

This hooks up call argument position completion to the typeCheckForCodeCompletion API to generate completions from all the solutions the constraint solver produces (even those requiring fixes), rather than relying on a single solution being applied to the AST (if any).

Co-authored-by: Nathan Hawes <nathan.john.hawes@gmail.com>

Author: ahoppen

include/swift/IDE/ArgumentCompletion.h

@@ -0,0 +1,79 @@
+//===--- ArgumentCompletion.h -----------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_IDE_ARGUMENTCOMPLETION_H
+#define SWIFT_IDE_ARGUMENTCOMPLETION_H
+
+#include "swift/IDE/CodeCompletionConsumer.h"
+#include "swift/IDE/CodeCompletionContext.h"
+#include "swift/IDE/PossibleParamInfo.h"
+#include "swift/Sema/CodeCompletionTypeChecking.h"
+
+namespace swift {
+namespace ide {
+
+class ArgumentTypeCheckCompletionCallback : public TypeCheckCompletionCallback {
+  struct Result {
+    /// The type associated with the code completion expression itself.
+    Type ExpectedType;
+    /// True if this is a subscript rather than a function call.
+    bool IsSubscript;
+    /// The FuncDecl or SubscriptDecl associated with the call.
+    ValueDecl *FuncD;
+    /// The type of the function being called.
+    Type FuncTy;
+    /// The index of the argument containing the completion location
+    unsigned ArgIdx;
+    /// The index of the parameter corresponding to the completion argument.
+    Optional<unsigned> ParamIdx;
+    /// The indices of all params that were bound to non-synthesized
+    /// arguments. Used so we don't suggest them even when the args are out of
+    /// order.
+    std::set<unsigned> ClaimedParamIndices;
+    /// True if the completion is a noninitial term in a variadic argument.
+    bool IsNoninitialVariadic;
+    /// The base type of the call/subscript (null for free functions).
+    Type BaseType;
+    /// True if an argument label precedes the completion location.
+    bool HasLabel;
+  };
+
+  CodeCompletionExpr *CompletionExpr;
+  SmallVector<Result, 4> Results;
+
+  /// Populates a vector of parameters to suggest along with a vector of types
+  /// to match the lookup results against.
+  ///
+  /// \Returns true if global lookup should be performed.
+  bool addPossibleParams(const ArgumentTypeCheckCompletionCallback::Result &Res,
+                         SmallVectorImpl<PossibleParamInfo> &Params,
+                         SmallVectorImpl<Type> &Types);
+
+public:
+  ArgumentTypeCheckCompletionCallback(CodeCompletionExpr *CompletionExpr)
+      : CompletionExpr(CompletionExpr) {}
+
+  void sawSolution(const constraints::Solution &solution) override;
+
+  /// \param IncludeSignature Whether to include a suggestion for the entire
+  /// function signature instead of suggesting individual labels. Used when
+  /// completing after the opening '(' of a function call \param Loc The
+  /// location of the code completion token
+  void deliverResults(bool IncludeSignature, SourceLoc Loc, DeclContext *DC,
+                      CodeCompletionContext &CompletionCtx,
+                      CodeCompletionConsumer &Consumer);
+};
+
+} // end namespace ide
+} // end namespace swift
+
+#endif // SWIFT_IDE_ARGUMENTCOMPLETION_H

include/swift/IDE/CodeCompletion.h

@@ -84,6 +84,10 @@ void lookupCodeCompletionResultsFromModule(CodeCompletionResultSink &targetSink,
                                            bool needLeadingDot,
                                            const SourceFile *SF);
 
+void addExprKeywords(CodeCompletionResultSink &Sink, DeclContext *DC);
+
+void addSuperKeyword(CodeCompletionResultSink &Sink, DeclContext *DC);
+
 } // end namespace ide
 } // end namespace swift
 

include/swift/IDE/PossibleParamInfo.h

@@ -33,6 +33,21 @@ struct PossibleParamInfo {
     assert((Param || !IsRequired) &&
            "nullptr with required flag is not allowed");
   };
+
+  friend bool operator==(const PossibleParamInfo &lhs,
+                         const PossibleParamInfo &rhs) {
+    bool ParamsMatch;
+    if (lhs.Param == nullptr && rhs.Param == nullptr) {
+      ParamsMatch = true;
+    } else if (lhs.Param == nullptr || rhs.Param == nullptr) {
+      // One is nullptr but the other is not.
+      ParamsMatch = false;
+    } else {
+      // Both are not nullptr.
+      ParamsMatch = (*lhs.Param == *rhs.Param);
+    }
+    return ParamsMatch && (lhs.IsRequired == rhs.IsRequired);
+  }
 };
 
 } // end namespace ide

include/swift/Sema/ConstraintSystem.h

@@ -2406,6 +2406,10 @@ class ConstraintSystem {
   /// diagnostics when result builder has multiple overloads.
   llvm::SmallDenseSet<AnyFunctionRef> InvalidResultBuilderBodies;
 
+  /// Arguments after the code completion token that were thus ignored (i.e.
+  /// assigned fresh type variables) for type checking.
+  llvm::SetVector<ConstraintLocator *> IgnoredArguments;
+
   /// Maps node types used within all portions of the constraint
   /// system, instead of directly using the types on the
   /// nodes themselves. This allows us to typecheck and
@@ -3171,9 +3175,24 @@ class ConstraintSystem {
     return TypeVariables.count(typeVar) > 0;
   }
 
-  /// Whether the given expression's source range contains the code
+  /// Whether the given ASTNode's source range contains the code
   /// completion location.
-  bool containsCodeCompletionLoc(Expr *expr) const;
+  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);
+  }
+
+  /// 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;
+  }
 
   void setClosureType(const ClosureExpr *closure, FunctionType *type) {
     assert(closure);
@@ -5534,8 +5553,44 @@ class MatchCallArgumentListener {
   /// \returns true to indicate that this should cause a failure, false
   /// otherwise.
   virtual bool relabelArguments(ArrayRef<Identifier> newNames);
+
+  /// \returns true if matchCallArguments should try to claim the argument at
+  /// \p argIndex while recovering from a failure. This is used to prevent
+  /// claiming of arguments after the code completion token.
+  virtual bool shouldClaimArgDuringRecovery(unsigned argIdx);
+
+  /// \returns true if \p arg can be claimed even though its argument label
+  /// doesn't match. This is the case for arguments representing the code
+  /// completion token if they don't contain a label. In these cases completion
+  /// will suggest the label.
+  virtual bool
+  canClaimArgIgnoringNameMismatch(const AnyFunctionType::Param &arg);
+};
+
+/// For a callsite containing a code completion expression, stores the index of
+/// the arg containing it along with the index of the first trailing closure and
+/// how many arguments were passed in total.
+struct CompletionArgInfo {
+  unsigned completionIdx;
+  Optional<unsigned> firstTrailingIdx;
+  unsigned argCount;
+
+  /// \returns true if the given argument index is possibly about to be written
+  /// by the user (given the completion index) so shouldn't be penalised as
+  /// missing when ranking solutions.
+  bool allowsMissingArgAt(unsigned argInsertIdx, AnyFunctionType::Param param);
+
+  /// \returns true if the argument containing the completion location is before
+  /// the argument with the given index.
+  bool isBefore(unsigned argIdx) { return completionIdx < argIdx; }
 };
 
+/// Extracts the index of the argument containing the code completion location
+/// from the provided anchor if it's a \c CallExpr, \c SubscriptExpr, or
+/// \c ObjectLiteralExpr.
+Optional<CompletionArgInfo> getCompletionArgInfo(ASTNode anchor,
+                                                 ConstraintSystem &cs);
+
 /// Match the call arguments (as described by the given argument type) to
 /// the parameters (as described by the given parameter type).
 ///