Merge pull request #32391 from xedin/type-check-for-code-completion

[TypeChecker] Add a entry point to be used for code completion

Author: xedin

lib/Sema/CSSolver.cpp

@@ -1375,6 +1375,49 @@ void ConstraintSystem::solveImpl(SmallVectorImpl<Solution> &solutions) {
   }
 }
 
+void ConstraintSystem::solveForCodeCompletion(
+    Expr *expr, DeclContext *DC, Type contextualType, ContextualTypePurpose CTP,
+    llvm::function_ref<void(const Solution &)> callback) {
+  // First, pre-check the expression, validating any types that occur in the
+  // expression and folding sequence expressions.
+  if (ConstraintSystem::preCheckExpression(expr, DC))
+    return;
+
+  ConstraintSystemOptions options;
+  options |= ConstraintSystemFlags::AllowFixes;
+  options |= ConstraintSystemFlags::SuppressDiagnostics;
+
+  ConstraintSystem cs(DC, options);
+
+  if (CTP != ContextualTypePurpose::CTP_Unused)
+    cs.setContextualType(expr, TypeLoc::withoutLoc(contextualType), CTP);
+
+  // Set up the expression type checker timer.
+  cs.Timer.emplace(expr, cs);
+
+  cs.shrink(expr);
+
+  if (cs.generateConstraints(expr, DC))
+    return;
+
+  llvm::SmallVector<Solution, 4> solutions;
+
+  {
+    SolverState state(cs, FreeTypeVariableBinding::Disallow);
+
+    // Enable "diagnostic mode" by default, this means that
+    // solver would produce "fixed" solutions along side of
+    // valid ones, which helps code completion to rank choices.
+    state.recordFixes = true;
+
+    cs.solveImpl(solutions);
+  }
+
+  for (const auto &solution : solutions) {
+    callback(solution);
+  }
+}
+
 void ConstraintSystem::collectDisjunctions(
     SmallVectorImpl<Constraint *> &disjunctions) {
   for (auto &constraint : InactiveConstraints) {

lib/Sema/ConstraintSystem.h

@@ -4728,6 +4728,31 @@ class ConstraintSystem {
                                  = FreeTypeVariableBinding::Disallow,
                                  bool allowFixes = false);
 
+  /// Construct and solve a system of constraints based on the given expression
+  /// and its contextual information.
+  ///
+  /// This menthod is designed to be used for code completion which means that
+  /// it doesn't mutate given expression, even if there is a single valid
+  /// solution, and constraint solver is allowed to produce partially correct
+  /// solutions, such solutions can have any number of holes in them, alongside
+  /// with valid ones.
+  ///
+  /// \param expr The expression involved in code completion.
+  ///
+  /// \param DC The declaration context this expression is found in.
+  ///
+  /// \param contextualType The type expression is being converted to.
+  ///
+  /// \param CTP When contextualType is specified, this indicates what
+  /// the conversion is doing.
+  ///
+  /// \param callback The callback to be used to provide results to
+  /// code completion.
+  static void
+  solveForCodeCompletion(Expr *expr, DeclContext *DC, Type contextualType,
+                         ContextualTypePurpose CTP,
+                         llvm::function_ref<void(const Solution &)> callback);
+
 private:
   /// Solve the system of constraints.
   ///

lib/Sema/TypeCheckConstraints.cpp

@@ -2328,6 +2328,19 @@ TypeChecker::getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
   }
 }
 
+void TypeChecker::typeCheckForCodeCompletion(
+    Expr *expr, DeclContext *DC, Type contextualType, ContextualTypePurpose CTP,
+    llvm::function_ref<void(const Solution &)> callback) {
+  auto &Context = DC->getASTContext();
+
+  FrontendStatsTracer StatsTracer(Context.Stats,
+                                  "typecheck-for-code-completion", expr);
+  PrettyStackTraceExpr stackTrace(Context, "code-completion", expr);
+
+  ConstraintSystem::solveForCodeCompletion(expr, DC, contextualType, CTP,
+                                           callback);
+}
+
 bool TypeChecker::typeCheckBinding(
     Pattern *&pattern, Expr *&initializer, DeclContext *DC,
     Type patternType, PatternBindingDecl *PBD, unsigned patternNumber) {

lib/Sema/TypeChecker.h

@@ -700,6 +700,17 @@ FunctionType *getTypeOfCompletionOperator(DeclContext *DC, Expr *LHS,
                                           DeclRefKind refKind,
                                           ConcreteDeclRef &refdDecl);
 
+/// Type check the given expression and provide results back to code completion
+/// via specified callback.
+///
+/// This menthod is designed to be used for code completion which means that
+/// it doesn't mutate AST and constraint solver is allowed to produce partially
+/// correct solutions, such solutions can have any number of holes in them,
+/// alongside with valid ones.
+void typeCheckForCodeCompletion(
+    Expr *expr, DeclContext *DC, Type contextualType, ContextualTypePurpose CTP,
+    llvm::function_ref<void(const constraints::Solution &)> callback);
+
 /// Check the key-path expression.
 ///
 /// Returns the type of the last component of the key-path.