[ConstraintSystem] Prevent shrink from solving "too complex" sub-expressions

Add additional checking for complexity of the shrinking candidate
given the number of the expressions solved so far and total number
of disjunctions present. This allows us to bail quicker in complex
expression cases which, at the very least, produces an error instead
of being "stuck" in solver for a long time.

Resolves: rdar://problem/32034560

Author: xedin

lib/Sema/CSSolver.cpp

@@ -1469,7 +1469,8 @@ ConstraintSystem::solveSingle(FreeTypeVariableBinding allowFreeTypeVariables) {
   return std::move(solutions[0]);
 }
 
-bool ConstraintSystem::Candidate::solve() {
+bool ConstraintSystem::Candidate::solve(
+    llvm::SmallDenseSet<Expr *> &shrunkExprs) {
   // Don't attempt to solve candidate if there is closure
   // expression involved, because it's handled specially
   // by parent constraint system (e.g. parameter lists).
@@ -1512,6 +1513,11 @@ bool ConstraintSystem::Candidate::solve() {
     return true;
   }
 
+  // If this candidate is too complex given the number
+  // of the domains we have reduced so far, let's bail out early.
+  if (isTooComplexGiven(&cs, shrunkExprs))
+    return false;
+
   if (TC.getLangOpts().DebugConstraintSolver) {
     auto &log = cs.getASTContext().TypeCheckerDebug->getStream();
     log << "--- Solving candidate for shrinking at ";
@@ -1565,7 +1571,7 @@ bool ConstraintSystem::Candidate::solve() {
   }
 
   // Record found solutions as suggestions.
-  this->applySolutions(solutions);
+  this->applySolutions(solutions, shrunkExprs);
 
   // Let's double-check if we have any implicit expressions
   // with type variables and nullify their types.
@@ -1577,7 +1583,8 @@ bool ConstraintSystem::Candidate::solve() {
 }
 
 void ConstraintSystem::Candidate::applySolutions(
-                            llvm::SmallVectorImpl<Solution> &solutions) const {
+    llvm::SmallVectorImpl<Solution> &solutions,
+    llvm::SmallDenseSet<Expr *> &shrunkExprs) 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.
@@ -1624,6 +1631,9 @@ void ConstraintSystem::Candidate::applySolutions(
       = TC.Context.AllocateUninitialized<ValueDecl *>(choices.size());
     std::uninitialized_copy(choices.begin(), choices.end(), decls.begin());
     OSR->setDecls(decls);
+
+    // Record successfully shrunk expression.
+    shrunkExprs.insert(OSR);
   }
 }
 
@@ -1692,7 +1702,8 @@ void ConstraintSystem::shrink(Expr *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_back({applyExpr, isa<OverloadSetRefExpr>(func)});
+        ApplyExprs.push_back(
+            {applyExpr, isa<OverloadSetRefExpr>(func) || isa<TypeExpr>(func)});
       }
 
       return { true, expr };
@@ -1926,11 +1937,12 @@ void ConstraintSystem::shrink(Expr *expr) {
   // so we can start solving them separately.
   expr->walk(collector);
 
+  llvm::SmallDenseSet<Expr *> shrunkExprs;
   for (auto &candidate : collector.Candidates) {
     // 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()) {
+    if (candidate.solve(shrunkExprs)) {
       // Let's restore all of the original OSR domains for this sub-expression,
       // this means that we can still make forward progress with solving of the
       // top sub-expressions.
@@ -1941,6 +1953,7 @@ void ConstraintSystem::shrink(Expr *expr) {
             return childExpr;
 
           OSR->setDecls(domain->getSecond());
+          shrunkExprs.erase(OSR);
         }
 
         return childExpr;

lib/Sema/ConstraintSystem.h

@@ -998,15 +998,45 @@ class ConstraintSystem {
     /// \brief Try to solve this candidate sub-expression
     /// and re-write it's OSR domains afterwards.
     ///
+    /// \param shrunkExprs The set of expressions which
+    /// domains have been successfully shrunk so far.
+    ///
     /// \returns true on solver failure, false otherwise.
-    bool solve();
+    bool solve(llvm::SmallDenseSet<Expr *> &shrunkExprs);
 
     /// \brief Apply solutions found by solver as reduced OSR sets for
     /// for current and all of it's sub-expressions.
     ///
     /// \param solutions The solutions found by running solver on the
     /// this candidate expression.
-    void applySolutions(llvm::SmallVectorImpl<Solution> &solutions) const;
+    ///
+    /// \param shrunkExprs The set of expressions which
+    /// domains have been successfully shrunk so far.
+    void applySolutions(llvm::SmallVectorImpl<Solution> &solutions,
+                        llvm::SmallDenseSet<Expr *> &shrunkExprs) const;
+
+    /// Check if attempt at solving of the candidate makes sense given
+    /// the current conditions - number of shrunk domains which is related
+    /// to the given candidate over the total number of disjunctions present.
+    static bool isTooComplexGiven(ConstraintSystem *const cs,
+                                  llvm::SmallDenseSet<Expr *> &shrunkExprs) {
+      SmallVector<Constraint *, 8> disjunctions;
+      cs->collectDisjunctions(disjunctions);
+
+      unsigned unsolvedDisjunctions = disjunctions.size();
+      for (auto *disjunction : disjunctions) {
+        auto *locator = disjunction->getLocator();
+        if (!locator)
+          continue;
+
+        if (auto *anchor = locator->getAnchor()) {
+          if (shrunkExprs.count(anchor) > 0)
+            --unsolvedDisjunctions;
+        }
+      }
+
+      return unsolvedDisjunctions >= 5;
+    }
   };
 
   /// \brief Describes the current solver state.

test/Sema/complex_expressions.swift

@@ -117,3 +117,12 @@ let sr3668Dict2: [Int: (Int, Int) -> Bool] =
      8: { $0 != $1 },  9: { $0 != $1 }, 10: { $0 != $1 }, 11: { $0 != $1 },
     12: { $0 != $1 }, 13: { $0 != $1 }, 14: { $0 != $1 }, 15: { $0 != $1 },
     16: { $0 != $1 }, 17: { $0 != $1 }, 18: { $0 != $1 }, 19: { $0 != $1 } ]
+
+// rdar://problem/32034560 - type-checker hangs trying to solve expression
+struct R32034560 {
+  private var R32034560: Array<Array<UInt32>>
+  private func foo(x: UInt32) -> UInt32 {
+    return ((self.R32034560[0][Int(x >> 24) & 0xFF] &+ self.R32034560[1][Int(x >> 16) & 0xFF]) ^ self.R32034560[2][Int(x >> 8) & 0xFF]) &+ self.R32034560[3][Int(x & 0xFF)]
+    // expected-error@-1 {{expression was too complex to be solved in reasonable time; consider breaking up the expression into distinct sub-expressions}}
+  }
+}