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.

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -236,36 +236,23 @@ struct RestrictionOrFix {
 
 
 class ExpressionTimer {
-public:
-  using AnchorType = llvm::PointerUnion<Expr *, ConstraintLocator *>;
-
-private:
-  AnchorType Anchor;
-  ASTContext &Context;
+  ConstraintSystem &CS;
   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,
-                  unsigned thresholdInSecs);
+  ExpressionTimer(ConstraintSystem &CS, unsigned thresholdInSecs);
 
   ~ExpressionTimer();
 
-  AnchorType getAnchor() const { return Anchor; }
-
-  SourceRange getAffectedRange() const;
-
-  unsigned getWarnLimit() const {
-    return Context.TypeCheckerOpts.WarnLongExpressionTypeChecking;
-  }
+  unsigned getWarnLimit() const;
   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 @@ class ConstraintSystem {
   /// \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 @@ class ConstraintSystem {
   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 @@ class ConstraintSystem {
   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 @@ class ConstraintSystem {
   /// 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 @@ class ConstraintSystem {
     return range.isValid() ? range : std::optional<SourceRange>();
   }
 
-  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(size_t solutionMemory);
 
   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)
-        .highlight(func->getBodySourceRange());
-      salvagedResult.markAsDiagnosed();
+      cs.diagnoseTooComplex(func->getLoc(), salvagedResult);
       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);
-
-  // Set up expression type checker timer for the candidate.
-  cs.startExpressionTimer(E);
+  cs.startExpression(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();
-
-      // 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();
+      diagnoseTooComplex(target.getLoc(), solution);
       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.
-    startExpressionTimer(expr);
+    startExpression(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>>
-      OuterTimeRemaining = std::nullopt;
+  std::optional<unsigned> OuterTimeRemaining = std::nullopt;
 
   /// Conjunction constraint associated with this step.
   Constraint *Conjunction;
@@ -910,7 +909,7 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
 
     if (cs.Timer) {
       auto remainingTime = cs.Timer->getRemainingProcessTimeInSeconds();
-      OuterTimeRemaining.emplace(cs.Timer->getAnchor(), remainingTime);
+      OuterTimeRemaining.emplace(remainingTime);
     }
   }
 
@@ -925,11 +924,8 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
     if (HadFailure)
       restoreBestScore();
 
-    if (OuterTimeRemaining) {
-      auto anchor = OuterTimeRemaining->first;
-      auto remainingTime = OuterTimeRemaining->second;
-      CS.Timer.emplace(anchor, CS, remainingTime);
-    }
+    if (OuterTimeRemaining)
+      CS.Timer.emplace(CS, *OuterTimeRemaining);
   }
 
   StepResult resume(bool prevFailed) override;