[Type checker] Introduce directional path consistency algorithm

DPC algorithm tries to solve individual sub-expressions and combine
resolved types as a way to reduce pre-existing OSR domains. Solving
is done bottom-up so each consecutive sub-expression tightens
possible solution domain even further.

(cherry picked from commit 8e04a84)

Author: xedin

include/swift/AST/Expr.h

@@ -1421,7 +1421,11 @@ class OverloadSetRefExpr : public Expr {
 
 public:
   ArrayRef<ValueDecl*> getDecls() const { return Decls; }
-  
+
+  void setDecls(ArrayRef<ValueDecl *> domain) {
+    Decls = domain;
+  }
+
   /// getBaseType - Determine the type of the base object provided for the
   /// given overload set, which is only non-null when dealing with an overloaded
   /// member reference.

lib/Sema/CSSolver.cpp

@@ -20,6 +20,8 @@
 #include "llvm/Support/SaveAndRestore.h"
 #include <memory>
 #include <tuple>
+#include <stack>
+#include <queue>
 using namespace swift;
 using namespace constraints;
 
@@ -1349,6 +1351,321 @@ ConstraintSystem::solveSingle(FreeTypeVariableBinding allowFreeTypeVariables) {
   return std::move(solutions[0]);
 }
 
+bool ConstraintSystem::Candidate::solve() {
+  // Cleanup after constraint system generation/solving,
+  // because it would assign types to expressions, which
+  // might interfere with solving higher-level expressions.
+  ExprCleaner cleaner(E);
+
+  // Allocate new constraint system for sub-expression.
+  ConstraintSystem cs(TC, DC, None);
+
+  // Set contextual type if present. This is done before constraint generation
+  // to give a "hint" to that operation about possible optimizations.
+  if (!CT.isNull())
+    cs.setContextualType(E, CT, CTP);
+
+  // Generate constraints for the new system.
+  if (auto generatedExpr = cs.generateConstraints(E)) {
+    E = generatedExpr;
+  } else {
+    // Failure to generate constraint system for sub-expression means we can't
+    // continue solving sub-expressions.
+    return true;
+  }
+
+  // If there is contextual type present, add an explicit "conversion"
+  // constraint to the system.
+  if (!CT.isNull()) {
+    auto constraintKind = ConstraintKind::Conversion;
+    if (CTP == CTP_CallArgument)
+      constraintKind = ConstraintKind::ArgumentConversion;
+
+    cs.addConstraint(constraintKind, E->getType(), CT.getType(),
+                     cs.getConstraintLocator(E), /*isFavored=*/true);
+  }
+
+  // Try to solve the system and record all available solutions.
+  llvm::SmallVector<Solution, 2> solutions;
+  {
+    SolverState state(cs);
+    cs.solverState = &state;
+
+    // Use solveRec() instead of solve() in here, because solve()
+    // would try to deduce the best solution, which we don't
+    // really want. Instead, we want the reduced set of domain choices.
+    cs.solveRec(solutions, FreeTypeVariableBinding::Allow);
+
+    cs.solverState = nullptr;
+  }
+
+  // No solutions for the sub-expression means that either main expression
+  // needs salvaging or it's inconsistent (read: doesn't have solutions).
+  if (solutions.empty())
+    return true;
+
+  // Record found solutions as suggestions.
+  this->applySolutions(solutions);
+  return false;
+}
+
+void ConstraintSystem::Candidate::applySolutions(
+                            llvm::SmallVectorImpl<Solution> &solutions) const {
+  // A collection of OSRs with their newly reduced domains,
+  // it's domains are sets because multiple solutions can have the same
+  // choice for one of the type variables, and we want no duplication.
+  llvm::SmallDenseMap<OverloadSetRefExpr *, llvm::SmallSet<ValueDecl *, 2>>
+    domains;
+  for (auto &solution : solutions) {
+    for (auto choice : solution.overloadChoices) {
+      // Some of the choices might not have locators.
+      if (!choice.getFirst())
+        continue;
+
+      auto anchor = choice.getFirst()->getAnchor();
+      // Anchor is not available or expression is not an overload set.
+      if (!anchor || !isa<OverloadSetRefExpr>(anchor))
+        continue;
+
+      auto OSR = cast<OverloadSetRefExpr>(anchor);
+      auto overload = choice.getSecond().choice;
+      auto type = overload.getDecl()->getInterfaceType();
+
+      // One of the solutions has polymorphic type assigned with one of it's
+      // type variables. Such functions can only be properly resolved
+      // via complete expression, so we'll have to forget solutions
+      // we have already recorded. They might not include all viable overload
+      // choices.
+      if (type->is<GenericFunctionType>()) {
+        return;
+      }
+
+      domains[OSR].insert(overload.getDecl());
+    }
+  }
+
+  // Reduce the domains.
+  for (auto &domain : domains) {
+    auto OSR = domain.getFirst();
+    auto &choices = domain.getSecond();
+
+    // If the domain wasn't reduced, skip it.
+    if (OSR->getDecls().size() == choices.size()) continue;
+
+    // Update the expression with the reduced domain.
+    MutableArrayRef<ValueDecl *> decls
+      = TC.Context.AllocateUninitialized<ValueDecl *>(choices.size());
+    std::uninitialized_copy(choices.begin(), choices.end(), decls.begin());
+    OSR->setDecls(decls);
+  }
+}
+
+void ConstraintSystem::shrink(Expr *expr) {
+  typedef llvm::SmallDenseMap<Expr *, ArrayRef<ValueDecl *>> DomainMap;
+
+  // A collection of original domains of all of the expressions,
+  // so they can be restored in case of failure.
+  DomainMap domains;
+
+  struct ExprCollector : public ASTWalker {
+    // The primary constraint system.
+    ConstraintSystem &CS;
+
+    // All of the sub-expressions of certain type (binary/unary/calls) in
+    // depth-first order.
+    std::queue<Candidate> &SubExprs;
+
+    // Counts the number of overload sets present in the tree so far.
+    // Note that the traversal is depth-first.
+    std::stack<std::pair<ApplyExpr *, unsigned>,
+               llvm::SmallVector<std::pair<ApplyExpr *, unsigned>, 4>>
+      ApplyExprs;
+
+    // A collection of original domains of all of the expressions,
+    // so they can be restored in case of failure.
+    DomainMap &Domains;
+
+    ExprCollector(ConstraintSystem &cs,
+                  std::queue<Candidate> &container,
+                  DomainMap &domains)
+      : CS(cs), SubExprs(container), Domains(domains) { }
+
+    std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
+      // A dictionary expression is just a set of tuples; try to solve ones
+      // that have overload sets.
+      if (auto dictionaryExpr = dyn_cast<DictionaryExpr>(expr)) {
+        for (auto element : dictionaryExpr->getElements()) {
+          unsigned numOverlaods = 0;
+          element->walk(OverloadSetCounter(numOverlaods));
+
+          // There are no overload sets in the element; skip it.
+          if (numOverlaods == 0)
+            continue;
+
+          // FIXME: Could we avoid creating a separate dictionary expression
+          // here by introducing a contextual type on the element?
+          auto dict = DictionaryExpr::create(CS.getASTContext(),
+                                             dictionaryExpr->getLBracketLoc(),
+                                             { element },
+                                             dictionaryExpr->getRBracketLoc(),
+                                             dictionaryExpr->getType());
+
+          // Make each of the dictionary elements an independent dictionary,
+          // such makes it easy to type-check everything separately.
+          SubExprs.push(Candidate(CS, dict));
+        }
+
+        // Don't try to walk into the dictionary.
+        return { false, expr };
+      }
+
+      // Let's not attempt to type-check closures or default values,
+      // which has already been type checked anyway.
+      if (isa<ClosureExpr>(expr) || isa<DefaultValueExpr>(expr)) {
+        return { false, expr };
+      }
+
+      // Coerce to type expressions are only viable if they have
+      // a single child expression.
+      if (auto coerceExpr = dyn_cast<CoerceExpr>(expr)) {
+        if (!coerceExpr->getSubExpr()) {
+          return { false, expr };
+        }
+      }
+
+      if (auto OSR = dyn_cast<OverloadSetRefExpr>(expr)) {
+        Domains[OSR] = OSR->getDecls();
+      }
+
+      if (auto applyExpr = dyn_cast<ApplyExpr>(expr)) {
+        auto func = applyExpr->getFn();
+        // Let's record this function application for post-processing
+        // as well as if it contains overload set, see walkToExprPost.
+        ApplyExprs.push({ applyExpr, isa<OverloadSetRefExpr>(func) });
+      }
+
+      return { true, expr };
+    }
+
+    Expr *walkToExprPost(Expr *expr) override {
+      if (!isa<ApplyExpr>(expr))
+        return expr;
+
+      unsigned numOverloadSets = 0;
+      // Let's count how many overload sets do we have.
+      while (!ApplyExprs.empty()) {
+        auto application = ApplyExprs.top();
+        auto applyExpr = application.first;
+
+        // Add overload sets tracked by current expression.
+        numOverloadSets += application.second;
+        ApplyExprs.pop();
+
+        // We've found the current expression, so record the number of
+        // overloads.
+        if (expr == applyExpr) {
+          ApplyExprs.push({ applyExpr, numOverloadSets });
+          break;
+        }
+      }
+
+      // If there are fewer than two overloads in the chain
+      // there is no point of solving this expression,
+      // because we won't be able to reduce it's domain.
+      if (numOverloadSets > 1)
+        SubExprs.push(Candidate(CS, expr));
+
+      return expr;
+    }
+  };
+
+  std::queue<Candidate> expressions;
+  ExprCollector collector(*this, expressions, domains);
+
+  // Collect all of the binary/unary and call sub-expressions
+  // so we can start solving them separately.
+  expr->walk(collector);
+
+  while (!expressions.empty()) {
+    auto &candidate = expressions.front();
+
+    // If there are no results, let's forget everything we know about the
+    // system so far. This actually is ok, because some of the expressions
+    // might require manual salvaging.
+    if (candidate.solve()) {
+      // Let's restore all of the original OSR domains.
+      for (auto &domain : domains) {
+        if (auto OSR = dyn_cast<OverloadSetRefExpr>(domain.getFirst())) {
+          OSR->setDecls(domain.getSecond());
+        }
+      }
+      break;
+    }
+
+    expressions.pop();
+  }
+}
+
+ConstraintSystem::SolutionKind
+ConstraintSystem::solve(Expr *&expr,
+                        Type convertType,
+                        ExprTypeCheckListener *listener,
+                        SmallVectorImpl<Solution> &solutions,
+                        FreeTypeVariableBinding allowFreeTypeVariables) {
+  assert(!solverState && "use solveRec for recursive calls");
+
+  // Try to shrink the system by reducing disjunction domains. This
+  // goes through every sub-expression and generate it's own sub-system, to
+  // try to reduce the domains of those subexpressions.
+  shrink(expr);
+
+  // Generate constraints for the main system.
+  if (auto generatedExpr = generateConstraints(expr))
+    expr = generatedExpr;
+  else {
+    return SolutionKind::Error;
+  }
+
+  // If there is a type that we're expected to convert to, add the conversion
+  // constraint.
+  if (convertType) {
+    auto constraintKind = ConstraintKind::Conversion;
+    if (getContextualTypePurpose() == CTP_CallArgument)
+      constraintKind = ConstraintKind::ArgumentConversion;
+
+    if (allowFreeTypeVariables == FreeTypeVariableBinding::UnresolvedType) {
+      convertType = convertType.transform([&](Type type) -> Type {
+        if (type->is<UnresolvedType>())
+          return createTypeVariable(getConstraintLocator(expr), 0);
+        return type;
+      });
+    }
+
+    addConstraint(constraintKind, expr->getType(), convertType,
+                  getConstraintLocator(expr), /*isFavored*/ true);
+  }
+
+  // Notify the listener that we've built the constraint system.
+  if (listener && listener->builtConstraints(*this, expr)) {
+    return SolutionKind::Error;
+  }
+
+  if (TC.getLangOpts().DebugConstraintSolver) {
+    auto &log = getASTContext().TypeCheckerDebug->getStream();
+    log << "---Initial constraints for the given expression---\n";
+    expr->print(log);
+    log << "\n";
+    print(log);
+  }
+
+  // Try to solve the constraint system using computed suggestions.
+  solve(solutions, allowFreeTypeVariables);
+
+  // If there are no solutions let's mark system as unsolved,
+  // and solved otherwise even if there are multiple solutions still present.
+  return solutions.empty() ? SolutionKind::Unsolved : SolutionKind::Solved;
+}
+
 bool ConstraintSystem::solve(SmallVectorImpl<Solution> &solutions,
                              FreeTypeVariableBinding allowFreeTypeVariables) {
   assert(!solverState && "use solveRec for recursive calls");
@@ -1372,7 +1689,7 @@ bool ConstraintSystem::solve(SmallVectorImpl<Solution> &solutions,
 
   // Remove the solver state.
   this->solverState = nullptr;
-  
+
   // We fail if there is no solution.
   return solutions.empty();
 }