Merge pull request #39373 from ahoppen/pr/solver-based-arg-completion

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

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

@@ -386,6 +386,10 @@ class TypeVariableType::Implementation {
 
   bool isTypeSequence() const;
 
+  /// Determine whether this type variable represents a code completion
+  /// expression.
+  bool isCodeCompletionToken() const;
+
   /// Retrieve the representative of the equivalence class to which this
   /// type variable belongs.
   ///
@@ -2409,6 +2413,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
@@ -3174,9 +3182,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);
@@ -5537,8 +5560,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).
 ///