[ConstraintSystem] Improve precision of "too complex" diagnostic

include/swift/Sema/ConstraintSystem.h

@@ -194,6 +194,8 @@ class ExpressionTimer {
 
   AnchorType getAnchor() const { return Anchor; }
 
+  SourceRange getAffectedRange() const;
+
   unsigned getWarnLimit() const {
     return Context.TypeCheckerOpts.WarnLongExpressionTypeChecking;
   }
@@ -5415,9 +5417,17 @@ class ConstraintSystem {
 
   /// Determine if we've already explored too many paths in an
   /// attempt to solve this expression.
-  bool isExpressionAlreadyTooComplex = false;
-  bool getExpressionTooComplex(size_t solutionMemory) {
-    if (isExpressionAlreadyTooComplex)
+  std::pair<bool, SourceRange> isAlreadyTooComplex = {false, SourceRange()};
+
+  /// If optional is not nil, result is guaranteed to point at a valid
+  /// location.
+  Optional<SourceRange> getTooComplexRange() const {
+    auto range = isAlreadyTooComplex.second;
+    return range.isValid() ? range : Optional<SourceRange>();
+  }
+
+  bool isTooComplex(size_t solutionMemory) {
+    if (isAlreadyTooComplex.first)
       return true;
 
     auto CancellationFlag = getASTContext().CancellationFlag;
@@ -5428,7 +5438,9 @@ class ConstraintSystem {
     MaxMemory = std::max(used, MaxMemory);
     auto threshold = getASTContext().TypeCheckerOpts.SolverMemoryThreshold;
     if (MaxMemory > threshold) {
-      return isExpressionAlreadyTooComplex= true;
+      // No particular location for OoM problems.
+      isAlreadyTooComplex.first = true;
+      return true;
     }
 
     if (Timer && Timer->isExpired()) {
@@ -5437,27 +5449,29 @@ class ConstraintSystem {
       // emitting an error.
       Timer->disableWarning();
 
-      return isExpressionAlreadyTooComplex = true;
+      isAlreadyTooComplex = {true, Timer->getAffectedRange()};
+      return true;
     }
 
     // Bail out once we've looked at a really large number of
     // choices.
     if (CountScopes > getASTContext().TypeCheckerOpts.SolverBindingThreshold) {
-      return isExpressionAlreadyTooComplex = true;
+      isAlreadyTooComplex.first = true;
+      return true;
     }
 
     return false;
   }
 
-  bool getExpressionTooComplex(SmallVectorImpl<Solution> const &solutions) {
-    if (isExpressionAlreadyTooComplex)
+  bool isTooComplex(SmallVectorImpl<Solution> const &solutions) {
+    if (isAlreadyTooComplex.first)
       return true;
 
     size_t solutionMemory = 0;
     for (auto const& s : solutions) {
       solutionMemory += s.getTotalMemory();
     }
-    return getExpressionTooComplex(solutionMemory);
+    return isTooComplex(solutionMemory);
   }
 
   // If the given constraint is an applied disjunction, get the argument function

include/swift/Sema/SolutionResult.h

@@ -64,6 +64,9 @@ class SolutionResult {
   /// \c numSolutions entries.
   Solution *solutions = nullptr;
 
+  /// A source range that was too complex to solve.
+  Optional<SourceRange> TooComplexAt = None;
+
   /// General constructor for the named constructors.
   SolutionResult(Kind kind) : kind(kind) {
     emittedDiagnostic = false;
@@ -95,9 +98,7 @@ class SolutionResult {
 
   /// Produce a "too complex" failure, which was not yet been
   /// diagnosed.
-  static SolutionResult forTooComplex() {
-    return SolutionResult(TooComplex);
-  }
+  static SolutionResult forTooComplex(Optional<SourceRange> affected);
 
   /// Produce a failure that has already been diagnosed.
   static SolutionResult forError() {
@@ -123,6 +124,10 @@ class SolutionResult {
   /// Take the set of solutions when there is an ambiguity.
   MutableArrayRef<Solution> takeAmbiguousSolutions() &&;
 
+  /// Retrieve a range of source that has been determined to be too
+  /// complex to solve in a reasonable time.
+  Optional<SourceRange> getTooComplexAt() const { return TooComplexAt; }
+
   /// Whether this solution requires the client to produce a diagnostic.
   bool requiresDiagnostic() const {
     switch (kind) {

lib/Sema/CSApply.cpp

@@ -9106,6 +9106,12 @@ SolutionResult SolutionResult::forAmbiguous(
   return result;
 }
 
+SolutionResult SolutionResult::forTooComplex(Optional<SourceRange> affected) {
+  SolutionResult result(Kind::TooComplex);
+  result.TooComplexAt = affected;
+  return result;
+}
+
 SolutionResult::~SolutionResult() {
   assert((!requiresDiagnostic() || emittedDiagnostic) &&
          "SolutionResult was destroyed without emitting a diagnostic");

lib/Sema/CSSolver.cpp

@@ -1298,12 +1298,21 @@ Optional<std::vector<Solution>> ConstraintSystem::solve(
       maybeProduceFallbackDiagnostic(target);
       return None;
 
-    case SolutionResult::TooComplex:
-      getASTContext().Diags.diagnose(
-          target.getLoc(), diag::expression_too_complex)
-            .highlight(target.getSourceRange());
+    case SolutionResult::TooComplex: {
+      auto affectedRange = solution.getTooComplexAt();
+
+      // If affected range is unknown, let's use whole
+      // target.
+      if (!affectedRange)
+        affectedRange = target.getSourceRange();
+
+      getASTContext()
+          .Diags.diagnose(affectedRange->Start, diag::expression_too_complex)
+          .highlight(*affectedRange);
+
       solution.markAsDiagnosed();
       return None;
+    }
 
     case SolutionResult::Ambiguous:
       // If salvaging produced an ambiguous result, it has already been
@@ -1374,8 +1383,8 @@ ConstraintSystem::solveImpl(SolutionApplicationTarget &target,
   SmallVector<Solution, 4> solutions;
   solve(solutions, allowFreeTypeVariables);
 
-  if (getExpressionTooComplex(solutions))
-    return SolutionResult::forTooComplex();
+  if (isTooComplex(solutions))
+    return SolutionResult::forTooComplex(getTooComplexRange());
 
   switch (solutions.size()) {
   case 0:
@@ -1415,7 +1424,7 @@ bool ConstraintSystem::solve(SmallVectorImpl<Solution> &solutions,
   filterSolutions(solutions);
 
   // We fail if there is no solution or the expression was too complex.
-  return solutions.empty() || getExpressionTooComplex(solutions);
+  return solutions.empty() || isTooComplex(solutions);
 }
 
 void ConstraintSystem::solveImpl(SmallVectorImpl<Solution> &solutions) {

lib/Sema/CSStep.cpp

@@ -251,7 +251,7 @@ bool SplitterStep::mergePartialSolutions() const {
 
     // Since merging partial solutions can go exponential, make sure we didn't
     // pass the "too complex" thresholds including allocated memory and time.
-    if (CS.getExpressionTooComplex(solutionMemory))
+    if (CS.isTooComplex(solutionMemory))
       return false;
 
   } while (nextCombination(counts, indices));
@@ -313,7 +313,7 @@ StepResult ComponentStep::take(bool prevFailed) {
   // One of the previous components created by "split"
   // failed, it means that we can't solve this component.
   if ((prevFailed && DependsOnPartialSolutions.empty()) ||
-      CS.getExpressionTooComplex(Solutions))
+      CS.isTooComplex(Solutions))
     return done(/*isSuccess=*/false);
 
   // Setup active scope, only if previous component didn't fail.

lib/Sema/CSStep.h

@@ -506,7 +506,7 @@ template <typename P> class BindingStep : public SolverStep {
   StepResult take(bool prevFailed) override {
     // Before attempting the next choice, let's check whether the constraint
     // system is too complex already.
-    if (CS.getExpressionTooComplex(Solutions))
+    if (CS.isTooComplex(Solutions))
       return done(/*isSuccess=*/false);
 
     while (auto choice = Producer()) {

lib/Sema/ConstraintSystem.cpp

@@ -56,6 +56,22 @@ ExpressionTimer::ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS,
       PrintDebugTiming(CS.getASTContext().TypeCheckerOpts.DebugTimeExpressions),
       PrintWarning(true) {}
 
+SourceRange ExpressionTimer::getAffectedRange() const {
+  ASTNode anchor;
+
+  if (auto *locator = Anchor.dyn_cast<ConstraintLocator *>()) {
+    anchor = simplifyLocatorToAnchor(locator);
+    // If locator couldn't be simplified down to a single AST
+    // element, let's use its root.
+    if (!anchor)
+      anchor = locator->getAnchor();
+  } else {
+    anchor = Anchor.get<Expr *>();
+  }
+
+  return anchor.getSourceRange();
+}
+
 ExpressionTimer::~ExpressionTimer() {
   auto elapsed = getElapsedProcessTimeInFractionalSeconds();
   unsigned elapsedMS = static_cast<unsigned>(elapsed * 1000);
@@ -81,22 +97,13 @@ ExpressionTimer::~ExpressionTimer() {
   if (WarnLimit == 0 || elapsedMS < WarnLimit)
     return;
 
-  ASTNode anchor;
-  if (auto *locator = Anchor.dyn_cast<ConstraintLocator *>()) {
-    anchor = simplifyLocatorToAnchor(locator);
-    // If locator couldn't be simplified down to a single AST
-    // element, let's warn about its root.
-    if (!anchor)
-      anchor = locator->getAnchor();
-  } else {
-    anchor = Anchor.get<Expr *>();
-  }
+  auto sourceRange = getAffectedRange();
 
-  if (anchor.getStartLoc().isValid()) {
+  if (sourceRange.Start.isValid()) {
     Context.Diags
-        .diagnose(anchor.getStartLoc(), diag::debug_long_expression, elapsedMS,
+        .diagnose(sourceRange.Start, diag::debug_long_expression, elapsedMS,
                   WarnLimit)
-        .highlight(anchor.getSourceRange());
+        .highlight(sourceRange);
   }
 }
 
@@ -3776,8 +3783,8 @@ SolutionResult ConstraintSystem::salvage() {
     // Fall through to produce diagnostics.
   }
 
-  if (getExpressionTooComplex(viable))
-    return SolutionResult::forTooComplex();
+  if (isTooComplex(viable))
+    return SolutionResult::forTooComplex(getTooComplexRange());
 
   // Could not produce a specific diagnostic; punt to the client.
   return SolutionResult::forUndiagnosedError();

validation-test/Sema/type_checker_perf/slow/rdar19612086.swift

@@ -11,12 +11,11 @@ struct rdar19612086 {
   let x = 1.0
 
   var description : String {
-    return "\(i)" + Stringly(format: "%.2f", x) +
+    return "\(i)" + Stringly(format: "%.2f", x) +   // expected-error {{reasonable time}}
            "\(i+1)" + Stringly(format: "%.2f", x) +
            "\(i+2)" + Stringly(format: "%.2f", x) +
            "\(i+3)" + Stringly(format: "%.2f", x) +
            "\(i+4)" + Stringly(format: "%.2f", x) +
            "\(i+5)" + Stringly(format: "%.2f", x)
-    // expected-error@-1 {{reasonable time}}
   }
 }

validation-test/Sema/type_checker_perf/slow/rdar22079400.swift

@@ -1,9 +1,8 @@
 // RUN: %target-typecheck-verify-swift -solver-expression-time-threshold=1
 // REQUIRES: tools-release,no_asan
 
-let _ = (0...1).lazy.flatMap {
+let _ = (0...1).lazy.flatMap { // expected-error {{reasonable time}}
   a in (1...2).lazy.map { b in (a, b) }
 }.filter {
-  // expected-error@-1 {{reasonable time}}
   1 < $0 && $0 < $1 && $0 + $1 < 3
 }

validation-test/Sema/type_checker_perf/slow/rdar91310777.swift

@@ -0,0 +1,16 @@
+// RUN: %target-typecheck-verify-swift -solver-expression-time-threshold=1
+// REQUIRES: tools-release,no_asan
+
+func compute(_ cont: () -> Void) {}
+
+func test() {
+  compute {
+    let x = 42
+    compute {
+      print(x)
+      let v: UInt64 = UInt64((24 / UInt32(1)) + UInt32(0) - UInt32(0) - 24 / 42 - 42)
+      // expected-error@-1 {{the compiler is unable to type-check this expression in reasonable time; try breaking up the expression into distinct sub-expressions}}
+      print(v)
+    }
+  }
+}