Sema: Fix source location bookkeeping for 'reasonable time' diagnostic

We already had bookkeeping to track which statement in a multi-statement
closure we were looking at, but this was only used for the 'reasonable time'
diagnostic in the case that we hit the expression timer, which was almost
never hit, and is now off by default. The scope, memory, and trial limits
couldn't use this information, so they would always diagnose the entire
target being type checked.

Move it up from ExpressionTimer to ConstraintSystem, so that we get the
right source location there too. Also, factor out some code duplication
in BuilderTransform to ensure we get the same benefit for result builders
applied to function bodies too.

include/swift/Sema/ConstraintSystem.h

@@ -236,36 +236,23 @@ public:
 
 
 class ExpressionTimer {
-public:
+  ConstraintSystem &CS;
-  using AnchorType = llvm::PointerUnion<Expr *, ConstraintLocator *>;
-
-private:
-  AnchorType Anchor;
-  ASTContext &Context;
   llvm::TimeRecord StartTime;
 
   /// The number of seconds from creation until
   /// this timer is considered expired.
   unsigned ThresholdInSecs;
 
-  bool PrintDebugTiming;
   bool PrintWarning;
 
 public:
   const static unsigned NoLimit = (unsigned) -1;
 
-  ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS,
+  ExpressionTimer(ConstraintSystem &CS, unsigned thresholdInSecs);
-                  unsigned thresholdInSecs);
 
   ~ExpressionTimer();
 
-  AnchorType getAnchor() const { return Anchor; }
+  unsigned getWarnLimit() const;
-
-  SourceRange getAffectedRange() const;
-
-  unsigned getWarnLimit() const {
-    return Context.TypeCheckerOpts.WarnLongExpressionTypeChecking;
-  }
   llvm::TimeRecord startedAt() const { return StartTime; }
 
   /// Return the elapsed process time (including fractional seconds)
@@ -2159,7 +2146,9 @@ struct ClosureIsolatedByPreconcurrency {
 /// solution of which assigns concrete types to each of the type variables.
 /// Constraint systems are typically generated given an (untyped) expression.
 class ConstraintSystem {
+private:
   ASTContext &Context;
+  SourceRange CurrentRange;
 
 public:
   DeclContext *DC;
@@ -5384,6 +5373,9 @@ private:
   /// \returns The selected conjunction.
   Constraint *selectConjunction();
 
+  void diagnoseTooComplex(SourceLoc fallbackLoc,
+                          SolutionResult &result);
+
   /// Solve the system of constraints generated from provided expression.
   ///
   /// \param target The target to generate constraints from.
@@ -5481,6 +5473,8 @@ private:
   compareSolutions(ConstraintSystem &cs, ArrayRef<Solution> solutions,
                    const SolutionDiff &diff, unsigned idx1, unsigned idx2);
 
+  void startExpressionTimer();
+
 public:
   /// Increase the score of the given kind for the current (partial) solution
   /// along the current solver path.
@@ -5518,7 +5512,6 @@ public:
   std::optional<unsigned> findBestSolution(SmallVectorImpl<Solution> &solutions,
                                            bool minimize);
 
-public:
   /// Apply a given solution to the target, producing a fully
   /// type-checked target or \c None if an error occurred.
   ///
@@ -5571,7 +5564,14 @@ public:
   /// resolved before any others.
   void optimizeConstraints(Expr *e);
 
-  void startExpressionTimer(ExpressionTimer::AnchorType anchor);
+  /// Set the current sub-expression (of a multi-statement closure, etc) for
+  /// the purposes of diagnosing "reasonable time" errors.
+  void startExpression(ASTNode node);
+
+  /// The source range of the target being type checked.
+  SourceRange getCurrentSourceRange() const {
+    return CurrentRange;
+  }
 
   /// Determine if we've already explored too many paths in an
   /// attempt to solve this expression.
@@ -5584,53 +5584,7 @@ public:
     return range.isValid() ? range : std::optional<SourceRange>();
   }
 
-  bool isTooComplex(size_t solutionMemory) {
+  bool isTooComplex(size_t solutionMemory);
-    if (isAlreadyTooComplex.first)
-      return true;
-
-    auto CancellationFlag = getASTContext().CancellationFlag;
-    if (CancellationFlag && CancellationFlag->load(std::memory_order_relaxed))
-      return true;
-
-    auto markTooComplex = [&](SourceRange range, StringRef reason) {
-      if (isDebugMode()) {
-        if (solverState)
-          llvm::errs().indent(solverState->getCurrentIndent());
-        llvm::errs() << "(too complex: " << reason << ")\n";
-      }
-      isAlreadyTooComplex = {true, range};
-      return true;
-    };
-
-    auto used = getASTContext().getSolverMemory() + solutionMemory;
-    MaxMemory = std::max(used, MaxMemory);
-    auto threshold = getASTContext().TypeCheckerOpts.SolverMemoryThreshold;
-    if (MaxMemory > threshold) {
-      // No particular location for OoM problems.
-      return markTooComplex(SourceRange(), "exceeded memory limit");
-    }
-
-    if (Timer && Timer->isExpired()) {
-      // Disable warnings about expressions that go over the warning
-      // threshold since we're arbitrarily ending evaluation and
-      // emitting an error.
-      Timer->disableWarning();
-
-      return markTooComplex(Timer->getAffectedRange(), "exceeded time limit");
-    }
-
-    auto &opts = getASTContext().TypeCheckerOpts;
-
-    // Bail out once we've looked at a really large number of choices.
-    if (opts.SolverScopeThreshold && NumSolverScopes > opts.SolverScopeThreshold)
-      return markTooComplex(SourceRange(), "exceeded scope limit");
-
-    // Bail out once we've taken a really large number of steps.
-    if (opts.SolverTrailThreshold && NumTrailSteps > opts.SolverTrailThreshold)
-      return markTooComplex(SourceRange(), "exceeded trail limit");
-
-    return false;
-  }
 
   bool isTooComplex(ArrayRef<Solution> solutions) {
     if (isAlreadyTooComplex.first)

lib/Sema/BuilderTransform.cpp

@@ -1053,9 +1053,7 @@ TypeChecker::applyResultBuilderBodyTransform(FuncDecl *func, Type builderType) {
 
     case SolutionResult::Kind::TooComplex:
       reportSolutionsToSolutionCallback(salvagedResult);
-      func->diagnose(diag::expression_too_complex)
+      cs.diagnoseTooComplex(func->getLoc(), salvagedResult);
-        .highlight(func->getBodySourceRange());
-      salvagedResult.markAsDiagnosed();
       return nullptr;
     }
 

lib/Sema/CSSolver.cpp

@@ -829,9 +829,7 @@ bool ConstraintSystem::Candidate::solve(
 
   // Allocate new constraint system for sub-expression.
   ConstraintSystem cs(DC, std::nullopt);
-
+  cs.startExpression(E);
-  // Set up expression type checker timer for the candidate.
-  cs.startExpressionTimer(E);
 
   // Generate constraints for the new system.
   if (auto generatedExpr = cs.generateConstraints(E, DC)) {
@@ -1455,18 +1453,7 @@ ConstraintSystem::solve(SyntacticElementTarget &target,
       return std::nullopt;
 
     case SolutionResult::TooComplex: {
-      auto affectedRange = solution.getTooComplexAt();
+      diagnoseTooComplex(target.getLoc(), solution);
-
-      // 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 std::nullopt;
     }
 
@@ -1501,6 +1488,19 @@ ConstraintSystem::solve(SyntacticElementTarget &target,
   llvm_unreachable("Loop always returns");
 }
 
+void ConstraintSystem::diagnoseTooComplex(SourceLoc fallbackLoc,
+                                          SolutionResult &result) {
+  auto affectedRange = result.getTooComplexAt();
+
+  SourceLoc loc = (affectedRange ? affectedRange->Start : fallbackLoc);
+  auto diag = getASTContext().Diags.diagnose(loc, diag::expression_too_complex);
+
+  if (affectedRange)
+    diag.highlight(*affectedRange);
+
+  result.markAsDiagnosed();
+}
+
 SolutionResult
 ConstraintSystem::solveImpl(SyntacticElementTarget &target,
                             FreeTypeVariableBinding allowFreeTypeVariables) {
@@ -1518,9 +1518,8 @@ ConstraintSystem::solveImpl(SyntacticElementTarget &target,
 
   assert(!solverState && "cannot be used directly");
 
-  // Set up the expression type checker timer.
   if (Expr *expr = target.getAsExpr())
-    startExpressionTimer(expr);
+    startExpression(expr);
 
   if (generateConstraints(target, allowFreeTypeVariables))
     return SolutionResult::forError();
@@ -1701,8 +1700,7 @@ bool ConstraintSystem::solveForCodeCompletion(
     // Tell the constraint system what the contextual type is.
     setContextualInfo(expr, target.getExprContextualTypeInfo());
 
-    // Set up the expression type checker timer.
+    startExpression(expr);
-    startExpressionTimer(expr);
 
     shrink(expr);
   }

lib/Sema/CSStep.cpp

@@ -823,9 +823,13 @@ bool ConjunctionStep::attempt(const ConjunctionElement &element) {
   // (expression) gets a fresh time slice to get solved. This
   // is important for closures with large number of statements
   // in them.
-  if (CS.Timer) {
+  if (CS.Timer)
     CS.Timer.reset();
-    CS.startExpressionTimer(element.getLocator());
+
+  {
+    auto *locator = element.getLocator();
+    auto anchor = simplifyLocatorToAnchor(locator);
+    CS.startExpression(anchor ? anchor : locator->getAnchor());
   }
 
   auto success = element.attempt(CS);

lib/Sema/CSStep.h

@@ -867,8 +867,7 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
 
   /// The number of milliseconds until outer constraint system
   /// is considered "too complex" if timer is enabled.
-  std::optional<std::pair<ExpressionTimer::AnchorType, unsigned>>
+  std::optional<unsigned> OuterTimeRemaining = std::nullopt;
-      OuterTimeRemaining = std::nullopt;
 
   /// Conjunction constraint associated with this step.
   Constraint *Conjunction;
@@ -910,7 +909,7 @@ public:
 
     if (cs.Timer) {
       auto remainingTime = cs.Timer->getRemainingProcessTimeInSeconds();
-      OuterTimeRemaining.emplace(cs.Timer->getAnchor(), remainingTime);
+      OuterTimeRemaining.emplace(remainingTime);
     }
   }
 
@@ -925,11 +924,8 @@ public:
     if (HadFailure)
       restoreBestScore();
 
-    if (OuterTimeRemaining) {
+    if (OuterTimeRemaining)
-      auto anchor = OuterTimeRemaining->first;
+      CS.Timer.emplace(CS, *OuterTimeRemaining);
-      auto remainingTime = OuterTimeRemaining->second;
-      CS.Timer.emplace(anchor, CS, remainingTime);
-    }
   }
 
   StepResult resume(bool prevFailed) override;

lib/Sema/ConstraintSystem.cpp

@@ -54,44 +54,82 @@ using namespace inference;
 
 #define DEBUG_TYPE "ConstraintSystem"
 
-ExpressionTimer::ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS,
+void ConstraintSystem::startExpression(ASTNode node) {
-                                 unsigned thresholdInSecs)
+  CurrentRange = node.getSourceRange();
-    : Anchor(Anchor), Context(CS.getASTContext()),
-      StartTime(llvm::TimeRecord::getCurrentTime()),
-      ThresholdInSecs(thresholdInSecs),
-      PrintDebugTiming(CS.getASTContext().TypeCheckerOpts.DebugTimeExpressions),
-      PrintWarning(true) {}
 
-SourceRange ExpressionTimer::getAffectedRange() const {
+  startExpressionTimer();
-  ASTNode anchor;
+}
 
-  if (auto *locator = Anchor.dyn_cast<ConstraintLocator *>()) {
+bool ConstraintSystem::isTooComplex(size_t solutionMemory) {
-    anchor = simplifyLocatorToAnchor(locator);
+  if (isAlreadyTooComplex.first)
-    // If locator couldn't be simplified down to a single AST
+    return true;
-    // element, let's use its root.
+
-    if (!anchor)
+  auto CancellationFlag = getASTContext().CancellationFlag;
-      anchor = locator->getAnchor();
+  if (CancellationFlag && CancellationFlag->load(std::memory_order_relaxed))
-  } else {
+    return true;
-    anchor = cast<Expr *>(Anchor);
+
+  auto markTooComplex = [&](SourceRange range, StringRef reason) {
+    if (isDebugMode()) {
+      if (solverState)
+        llvm::errs().indent(solverState->getCurrentIndent());
+      llvm::errs() << "(too complex: " << reason << ")\n";
+    }
+    isAlreadyTooComplex = {true, range};
+    return true;
+  };
+
+  auto used = getASTContext().getSolverMemory() + solutionMemory;
+  MaxMemory = std::max(used, MaxMemory);
+  auto threshold = getASTContext().TypeCheckerOpts.SolverMemoryThreshold;
+  if (MaxMemory > threshold) {
+    // Bail once we've used too much constraint solver arena memory.
+    return markTooComplex(getCurrentSourceRange(), "exceeded memory limit");
   }
 
-  return anchor.getSourceRange();
+  if (Timer && Timer->isExpired()) {
+    // Disable warnings about expressions that go over the warning
+    // threshold since we're arbitrarily ending evaluation and
+    // emitting an error.
+    Timer->disableWarning();
+
+    return markTooComplex(getCurrentSourceRange(), "exceeded time limit");
+  }
+
+  auto &opts = getASTContext().TypeCheckerOpts;
+
+  // Bail out once we've looked at a really large number of choices.
+  if (opts.SolverScopeThreshold && NumSolverScopes > opts.SolverScopeThreshold)
+    return markTooComplex(getCurrentSourceRange(), "exceeded scope limit");
+
+  // Bail out once we've taken a really large number of steps.
+  if (opts.SolverTrailThreshold && NumTrailSteps > opts.SolverTrailThreshold)
+    return markTooComplex(getCurrentSourceRange(), "exceeded trail limit");
+
+  return false;
+}
+
+ExpressionTimer::ExpressionTimer(ConstraintSystem &CS, unsigned thresholdInSecs)
+    : CS(CS),
+      StartTime(llvm::TimeRecord::getCurrentTime()),
+      ThresholdInSecs(thresholdInSecs),
+      PrintWarning(true) {}
+
+unsigned ExpressionTimer::getWarnLimit() const {
+  return CS.getASTContext().TypeCheckerOpts.WarnLongExpressionTypeChecking;
 }
 
 ExpressionTimer::~ExpressionTimer() {
   auto elapsed = getElapsedProcessTimeInFractionalSeconds();
   unsigned elapsedMS = static_cast<unsigned>(elapsed * 1000);
+  auto &ctx = CS.getASTContext();
 
-  if (PrintDebugTiming) {
+  auto range = CS.getCurrentSourceRange();
+
+  if (ctx.TypeCheckerOpts.DebugTimeExpressions) {
     // Round up to the nearest 100th of a millisecond.
     llvm::errs() << llvm::format("%0.2f", std::ceil(elapsed * 100000) / 100)
                  << "ms\t";
-    if (auto *E = Anchor.dyn_cast<Expr *>()) {
+    range.Start.print(llvm::errs(), ctx.SourceMgr);
-      E->getLoc().print(llvm::errs(), Context.SourceMgr);
-    } else {
-      auto *locator = cast<ConstraintLocator *>(Anchor);
-      locator->dump(&Context.SourceMgr, llvm::errs());
-    }
     llvm::errs() << "\n";
   }
 
@@ -103,13 +141,11 @@ ExpressionTimer::~ExpressionTimer() {
   if (WarnLimit == 0 || elapsedMS < WarnLimit)
     return;
 
-  auto sourceRange = getAffectedRange();
+  if (range.Start.isValid()) {
-
+    ctx.Diags
-  if (sourceRange.Start.isValid()) {
+        .diagnose(range.Start, diag::debug_long_expression, elapsedMS,
-    Context.Diags
-        .diagnose(sourceRange.Start, diag::debug_long_expression, elapsedMS,
                   WarnLimit)
-        .highlight(sourceRange);
+        .highlight(range);
   }
 }
 
@@ -140,7 +176,7 @@ ConstraintSystem::~ConstraintSystem() {
   }
 }
 
-void ConstraintSystem::startExpressionTimer(ExpressionTimer::AnchorType anchor) {
+void ConstraintSystem::startExpressionTimer() {
   ASSERT(!Timer);
 
   const auto &opts = getASTContext().TypeCheckerOpts;
@@ -156,7 +192,7 @@ void ConstraintSystem::startExpressionTimer(ExpressionTimer::AnchorType anchor)
     timeout = ExpressionTimer::NoLimit;
   }
 
-  Timer.emplace(anchor, *this, timeout);
+  Timer.emplace(*this, timeout);
 }
 
 void ConstraintSystem::incrementScopeCounter() {

test/Constraints/too_complex_source_location.swift

@@ -0,0 +1,19 @@
+// RUN: %target-typecheck-verify-swift -solver-scope-threshold=10
+
+// Note: the scope threshold is intentionally set low so that the expression will fail.
+//
+// The purpose of the test is to ensure the diagnostic points at the second statement in
+// the closure, and not the closure itself.
+// 
+// If the expression becomes very fast and we manage to type check it with fewer than
+// 10 scopes, please *do not* remove the expected error! Instead, make the expression
+// more complex again.
+
+let s = ""
+let n = 0
+
+let closure = {
+  let _ = 0
+  let _ = "" + s + "" + s + "" + s + "" + n + "" // expected-error {{reasonable time}}
+  let _ = 0
+}

validation-test/Sema/SwiftUI/too_complex_source_location.swift

@@ -0,0 +1,43 @@
+// RUN: %target-typecheck-verify-swift -target %target-cpu-apple-macosx10.15 -swift-version 5
+// REQUIRES: objc_interop
+
+// https://forums.swift.org/t/roadmap-for-improving-the-type-checker/82952/9
+//
+// The purpose of the test is to ensure the diagnostic points at the right statement in
+// the function body, and not the function declaration itself.
+//
+// Ideally, we would produce a useful diagnostic here. Once we are able to do that, we
+// will need to devise a new test which complains with 'reasonable time' to ensure the
+// source location remains correct.
+
+import SwiftUI
+
+struct ContentView: View {
+    @State var selection = ""
+
+    @State var a: Int?
+    @State var b: Int?
+    @State var c: Int?
+
+    var body: some View {
+        ScrollView {
+            VStack {
+                Picker(selection: $selection) {
+                    ForEach(["a", "b", "c"], id: \.self) {
+                        Text($0)  // expected-error {{reasonable time}}
+                            .foregroundStyl(.red) // Typo is here
+                    }
+                } label: {
+                }
+                .pickerStyle(.segmented)
+            }
+            .padding(.horizontal)
+        }
+        .onChange(of: a) { oldValue, newValue in
+        }
+        .onChange(of: b) { oldValue, newValue in
+        }
+        .onChange(of: c) { oldValue, newValue in
+        }
+    }
+}

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

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