[Constraint solver] Generalize solve() for arbitrary solution targets.

Start cleaning up the main “solve” entry point for solving an expression
and applying the solution, so it handles arbitrary solution targets.
This is another small step that doesn’t do much on its own, but will help
with unifying the various places in the code base where we run the solver.

Author: DougGregor

lib/Sema/CSApply.cpp

@@ -7331,7 +7331,8 @@ Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
     // If we're supposed to convert the expression to some particular type,
     // do so now.
     if (shouldCoerceToContextualType()) {
-      resultExpr = rewriter.coerceToType(resultExpr, convertType,
+      resultExpr = rewriter.coerceToType(resultExpr,
+                                         simplifyType(convertType),
                                          getConstraintLocator(expr));
     } else if (getType(resultExpr)->hasLValueType() &&
                !target.isDiscardedExpr()) {

lib/Sema/CSSolver.cpp

@@ -1074,22 +1074,22 @@ void ConstraintSystem::shrink(Expr *expr) {
   }
 }
 
-static bool debugConstraintSolverForExpr(ASTContext &C, Expr *expr) {
+static bool debugConstraintSolverForTarget(
+   ASTContext &C, SolutionApplicationTarget target) {
   if (C.TypeCheckerOpts.DebugConstraintSolver)
     return true;
 
   if (C.TypeCheckerOpts.DebugConstraintSolverOnLines.empty())
     // No need to compute the line number to find out it's not present.
     return false;
 
-  // Get the lines on which the expression starts and ends.
+  // Get the lines on which the target starts and ends.
   unsigned startLine = 0, endLine = 0;
-  if (expr->getSourceRange().isValid()) {
-    auto range =
-      Lexer::getCharSourceRangeFromSourceRange(C.SourceMgr,
-                                               expr->getSourceRange());
-    startLine = C.SourceMgr.getLineNumber(range.getStart());
-    endLine = C.SourceMgr.getLineNumber(range.getEnd());
+  SourceRange range = target.getSourceRange();
+  if (range.isValid()) {
+    auto charRange = Lexer::getCharSourceRangeFromSourceRange(C.SourceMgr, range);
+    startLine = C.SourceMgr.getLineNumber(charRange.getStart());
+    endLine = C.SourceMgr.getLineNumber(charRange.getEnd());
   }
 
   assert(startLine <= endLine && "expr ends before it starts?");
@@ -1107,25 +1107,26 @@ static bool debugConstraintSolverForExpr(ASTContext &C, Expr *expr) {
   return startBound != endBound;
 }
 
-bool ConstraintSystem::solve(Expr *&expr,
-                             Type convertType,
-                             ExprTypeCheckListener *listener,
-                             SmallVectorImpl<Solution> &solutions,
-                             FreeTypeVariableBinding allowFreeTypeVariables) {
+Optional<std::vector<Solution>> ConstraintSystem::solve(
+    SolutionApplicationTarget &target,
+    ExprTypeCheckListener *listener,
+    FreeTypeVariableBinding allowFreeTypeVariables
+) {
   llvm::SaveAndRestore<bool> debugForExpr(
       getASTContext().TypeCheckerOpts.DebugConstraintSolver,
-      debugConstraintSolverForExpr(getASTContext(), expr));
+      debugConstraintSolverForTarget(getASTContext(), target));
 
   /// Dump solutions for debugging purposes.
-  auto dumpSolutions = [&] {
+  auto dumpSolutions = [&](const SolutionResult &result) {
     // Debug-print the set of solutions.
     if (getASTContext().TypeCheckerOpts.DebugConstraintSolver) {
       auto &log = getASTContext().TypeCheckerDebug->getStream();
-      if (solutions.size() == 1) {
+      if (result.getKind() == SolutionResult::Success) {
         log << "---Solution---\n";
-        solutions[0].dump(log);
-      } else {
-        for (unsigned i = 0, e = solutions.size(); i != e; ++i) {
+        result.getSolution().dump(log);
+      } else if (result.getKind() == SolutionResult::Ambiguous) {
+        auto solutions = result.getAmbiguousSolutions();
+        for (unsigned i : indices(solutions)) {
           log << "--- Solution #" << i << " ---\n";
           solutions[i].dump(log);
         }
@@ -1135,61 +1136,61 @@ bool ConstraintSystem::solve(Expr *&expr,
 
   // Take up to two attempts at solving the system. The first attempts to
   // solve a system that is expected to be well-formed, the second kicks in
-  // when there is an error and attempts to salvage an ill-formed expression.
-  SolutionApplicationTarget target(expr, convertType, /*isDiscarded=*/false);
+  // when there is an error and attempts to salvage an ill-formed program.
   for (unsigned stage = 0; stage != 2; ++stage) {
     auto solution = (stage == 0)
         ? solveImpl(target, listener, allowFreeTypeVariables)
         : salvage();
 
     switch (solution.getKind()) {
-    case SolutionResult::Success:
+    case SolutionResult::Success: {
       // Return the successful solution.
-      solutions.clear();
-      solutions.push_back(std::move(solution).takeSolution());
-      dumpSolutions();
-      return false;
+      dumpSolutions(solution);
+      std::vector<Solution> result;
+      result.push_back(std::move(solution).takeSolution());
+      return std::move(result);
+    }
 
     case SolutionResult::Error:
-      return true;
+      return None;
 
     case SolutionResult::TooComplex:
-      getASTContext().Diags.diagnose(expr->getLoc(), diag::expression_too_complex)
-          .highlight(expr->getSourceRange());
+      getASTContext().Diags.diagnose(
+          target.getLoc(), diag::expression_too_complex)
+            .highlight(target.getSourceRange());
       solution.markAsDiagnosed();
-      return true;
+      return None;
 
     case SolutionResult::Ambiguous:
       // If salvaging produced an ambiguous result, it has already been
       // diagnosed.
       if (stage == 1) {
         solution.markAsDiagnosed();
-        return true;
+        return None;
       }
 
       if (Options.contains(
             ConstraintSystemFlags::AllowUnresolvedTypeVariables)) {
+        dumpSolutions(solution);
         auto ambiguousSolutions = std::move(solution).takeAmbiguousSolutions();
-        solutions.assign(std::make_move_iterator(ambiguousSolutions.begin()),
-                         std::make_move_iterator(ambiguousSolutions.end()));
-        dumpSolutions();
-        return false;
+        std::vector<Solution> result(
+            std::make_move_iterator(ambiguousSolutions.begin()),
+            std::make_move_iterator(ambiguousSolutions.end()));
+        return std::move(result);
       }
 
       LLVM_FALLTHROUGH;
 
     case SolutionResult::UndiagnosedError:
       if (shouldSuppressDiagnostics()) {
         solution.markAsDiagnosed();
-        return true;
+        return None;
       }
 
       if (stage == 1) {
-        diagnoseFailureFor(
-            SolutionApplicationTarget(expr, convertType,
-                                      /*isDiscarded=*/false));
+        diagnoseFailureFor(target);
         solution.markAsDiagnosed();
-        return true;
+        return None;
       }
 
       // Loop again to try to salvage.

lib/Sema/ConstraintSystem.h

@@ -1194,6 +1194,11 @@ class SolutionApplicationTarget {
     return expression.convertType;
   }
 
+  void setExprConversionType(Type type) {
+    assert(kind == Kind::expression);
+    expression.convertType = type;
+  }
+  
   bool isDiscardedExpr() const {
     assert(kind == Kind::expression);
     return expression.isDiscarded;
@@ -1234,6 +1239,18 @@ class SolutionApplicationTarget {
       return function.body->getSourceRange();
     }
   }
+
+  /// Retrieve the source location for the target.
+  SourceLoc getLoc() const {
+    switch (kind) {
+    case Kind::expression:
+      return expression.expression->getLoc();
+
+    case Kind::function:
+      return function.function.getLoc();
+    }
+  }
+
   /// Walk the contents of the application target.
   SolutionApplicationTarget walk(ASTWalker &walker);
 };
@@ -4104,26 +4121,21 @@ class ConstraintSystem {
   static bool preCheckExpression(Expr *&expr, DeclContext *dc,
                                  ConstraintSystem *baseCS = nullptr);
 
-  /// Solve the system of constraints generated from provided expression.
-  ///
-  /// The expression should have already been pre-checked with
-  /// preCheckExpression().
+  /// Solve the system of constraints generated from provided target.
   ///
-  /// \param expr The expression to generate constraints from.
-  /// \param convertType The expected type of the expression.
+  /// \param target The target that we'll generate constraints from, which
+  /// may be updated by the solving process.
   /// \param listener The callback to check solving progress.
-  /// \param solutions The set of solutions to the system of constraints.
   /// \param allowFreeTypeVariables How to bind free type variables in
   /// the solution.
   ///
-  /// \returns true is an error occurred, false is system is consistent
-  /// and solutions were found.
-  bool solve(Expr *&expr,
-             Type convertType,
-             ExprTypeCheckListener *listener,
-             SmallVectorImpl<Solution> &solutions,
-             FreeTypeVariableBinding allowFreeTypeVariables
-             = FreeTypeVariableBinding::Disallow);
+  /// \returns the set of solutions, if any were found, or \c None if an
+  /// error occurred. When \c None, an error has been emitted.
+  Optional<std::vector<Solution>> solve(
+      SolutionApplicationTarget &target,
+      ExprTypeCheckListener *listener,
+      FreeTypeVariableBinding allowFreeTypeVariables
+        = FreeTypeVariableBinding::Disallow);
 
   /// Solve the system of constraints.
   ///

lib/Sema/TypeCheckConstraints.cpp

@@ -2138,34 +2138,36 @@ Type TypeChecker::typeCheckExpression(Expr *&expr, DeclContext *dc,
     convertTo = getOptionalType(expr->getLoc(), var);
   }
 
-  SmallVector<Solution, 4> viable;
   // Attempt to solve the constraint system.
-  if (cs.solve(expr, convertTo, listener, viable, allowFreeTypeVariables))
+  SolutionApplicationTarget target(
+      expr, convertTo,
+      options.contains(TypeCheckExprFlags::IsDiscarded));
+  auto viable = cs.solve(target, listener, allowFreeTypeVariables);
+  if (!viable)
     return Type();
 
   // If the client allows the solution to have unresolved type expressions,
   // check for them now.  We cannot apply the solution with unresolved TypeVars,
   // because they will leak out into arbitrary places in the resultant AST.
   if (options.contains(TypeCheckExprFlags::AllowUnresolvedTypeVariables) &&
-      (viable.size() != 1 ||
+       (viable->size() != 1 ||
        (convertType.getType() && convertType.getType()->hasUnresolvedType()))) {
     return ErrorType::get(Context);
   }
 
-  auto result = expr;
-  auto &solution = viable[0];
+  auto result = target.getAsExpr();
+  auto &solution = (*viable)[0];
   if (!result)
     return Type();
 
   // Apply this solution to the constraint system.
   cs.applySolution(solution);
 
   // Apply the solution to the expression.
-  SolutionApplicationTarget target(
-      result, convertType.getType(),
-      options.contains(TypeCheckExprFlags::IsDiscarded));
   bool performingDiagnostics =
       options.contains(TypeCheckExprFlags::SubExpressionDiagnostics);
+  // FIXME: HACK!
+  target.setExprConversionType(convertType.getType());
   auto resultTarget = cs.applySolution(solution, target, performingDiagnostics);
   if (!resultTarget) {
     // Failure already diagnosed, above, as part of applying the solution.
@@ -2244,7 +2246,6 @@ getTypeOfExpressionWithoutApplying(Expr *&expr, DeclContext *dc,
   ConstraintSystem cs(dc, ConstraintSystemFlags::SuppressDiagnostics);
 
   // Attempt to solve the constraint system.
-  SmallVector<Solution, 4> viable;
   const Type originalType = expr->getType();
   const bool needClearType = originalType && originalType->hasError();
   const auto recoverOriginalType = [&] () {
@@ -2256,14 +2257,16 @@ getTypeOfExpressionWithoutApplying(Expr *&expr, DeclContext *dc,
   // re-check.
   if (needClearType)
     expr->setType(Type());
-  if (cs.solve(expr, /*convertType*/Type(), listener, viable,
-               allowFreeTypeVariables)) {
+  SolutionApplicationTarget target(expr, Type(), /*isDiscarded=*/false);
+  auto viable = cs.solve(target, listener, allowFreeTypeVariables);
+  if (!viable) {
     recoverOriginalType();
     return Type();
   }
 
   // Get the expression's simplified type.
-  auto &solution = viable[0];
+  expr = target.getAsExpr();
+  auto &solution = (*viable)[0];
   auto &solutionCS = solution.getConstraintSystem();
   Type exprType = solution.simplifyType(solutionCS.getType(expr));
 
@@ -2330,19 +2333,18 @@ void TypeChecker::getPossibleTypesOfExpressionWithoutApplying(
   ConstraintSystem cs(dc, options);
 
   // Attempt to solve the constraint system.
-  SmallVector<Solution, 4> viable;
-
   const Type originalType = expr->getType();
   if (originalType && originalType->hasError())
     expr->setType(Type());
 
-  cs.solve(expr, /*convertType*/ Type(), listener, viable,
-           allowFreeTypeVariables);
-
-  for (auto &solution : viable) {
-    auto exprType = solution.simplifyType(cs.getType(expr));
-    assert(exprType && !exprType->hasTypeVariable());
-    types.insert(exprType.getPointer());
+  SolutionApplicationTarget target(expr, Type(), /*isDiscarded=*/false);
+  if (auto viable = cs.solve(target, listener, allowFreeTypeVariables)) {
+    expr = target.getAsExpr();
+    for (auto &solution : *viable) {
+      auto exprType = solution.simplifyType(cs.getType(expr));
+      assert(exprType && !exprType->hasTypeVariable());
+      types.insert(exprType.getPointer());
+    }
   }
 }