Merge pull request #41282 from xedin/too-complex-with-multistmt-cljs

[ConstraintSystem] Adjust expression complexity computation to account for multi-statement closures

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -167,18 +167,31 @@ struct RestrictionOrFix {
 
 
 class ExpressionTimer {
-  Expr* E;
+public:
+  using AnchorType = llvm::PointerUnion<Expr *, ConstraintLocator *>;
+
+private:
+  AnchorType Anchor;
   ASTContext &Context;
   llvm::TimeRecord StartTime;
 
+  /// The number of milliseconds from creation until
+  /// this timer is considered expired.
+  unsigned ThresholdInMillis;
+
   bool PrintDebugTiming;
   bool PrintWarning;
 
 public:
-  ExpressionTimer(Expr *E, ConstraintSystem &CS);
+  /// This constructor sets a default threshold defined for all expressions
+  /// via compiler flag `solver-expression-time-threshold`.
+  ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS);
+  ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS, unsigned thresholdInMillis);
 
   ~ExpressionTimer();
 
+  AnchorType getAnchor() const { return Anchor; }
+
   unsigned getWarnLimit() const {
     return Context.TypeCheckerOpts.WarnLongExpressionTypeChecking;
   }
@@ -192,13 +205,19 @@ class ExpressionTimer {
     return endTime.getProcessTime() - StartTime.getProcessTime();
   }
 
+  /// Return the remaining process time in milliseconds until the
+  /// threshold specified during construction is reached.
+  unsigned getRemainingProcessTimeInMillis() const {
+    auto elapsed = unsigned(getElapsedProcessTimeInFractionalSeconds());
+    return elapsed >= ThresholdInMillis ? 0 : ThresholdInMillis - elapsed;
+  }
+
   // Disable emission of warnings about expressions that take longer
   // than the warning threshold.
   void disableWarning() { PrintWarning = false; }
 
-  bool isExpired(unsigned thresholdInMillis) const {
-    auto elapsed = getElapsedProcessTimeInFractionalSeconds();
-    return unsigned(elapsed) > thresholdInMillis;
+  bool isExpired() const {
+    return getRemainingProcessTimeInMillis() == 0;
   }
 };
 
@@ -5244,9 +5263,7 @@ class ConstraintSystem {
       return isExpressionAlreadyTooComplex= true;
     }
 
-    const auto timeoutThresholdInMillis =
-        getASTContext().TypeCheckerOpts.ExpressionTimeoutThreshold;
-    if (Timer && Timer->isExpired(timeoutThresholdInMillis)) {
+    if (Timer && Timer->isExpired()) {
       // Disable warnings about expressions that go over the warning
       // threshold since we're arbitrarily ending evaluation and
       // emitting an error.
@@ -5687,6 +5704,8 @@ class ConjunctionElement {
 
   bool attempt(ConstraintSystem &cs) const;
 
+  ConstraintLocator *getLocator() const { return Element->getLocator(); }
+
   void print(llvm::raw_ostream &Out, SourceManager *SM) const {
     Out << "conjunction element ";
     Element->print(Out, SM);

lib/Sema/CSStep.cpp

@@ -816,6 +816,18 @@ bool ConjunctionStep::attempt(const ConjunctionElement &element) {
   // by dropping all scoring information.
   CS.CurrentScore = Score();
 
+  // Reset the scope counter to avoid "too complex" failures
+  // when closure has a lot of elements in the body.
+  CS.CountScopes = 0;
+
+  // If timer is enabled, let's reset it so that each 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) {
+    CS.Timer.emplace(element.getLocator(), CS);
+  }
+
   auto success = element.attempt(CS);
 
   // If element attempt has failed, mark whole conjunction

lib/Sema/CSStep.h

@@ -857,6 +857,15 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
   /// The score established before conjunction is attempted.
   Score CurrentScore;
 
+  /// The number of constraint solver scopes already explored
+  /// before accepting this conjunction.
+  llvm::SaveAndRestore<unsigned> OuterScopeCount;
+
+  /// The number of milliseconds until outer constraint system
+  /// is considered "too complex" if timer is enabled.
+  Optional<std::pair<ExpressionTimer::AnchorType, unsigned>>
+      OuterTimeRemaining = None;
+
   /// Conjunction constraint associated with this step.
   Constraint *Conjunction;
   /// Position of the conjunction in the inactive constraints
@@ -889,13 +898,18 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
       : BindingStep(cs, {cs, conjunction},
                     conjunction->isIsolated() ? IsolatedSolutions : solutions),
         BestScore(getBestScore()), CurrentScore(getCurrentScore()),
-        Conjunction(conjunction), AfterConjunction(erase(conjunction)),
-        OuterSolutions(solutions) {
+        OuterScopeCount(cs.CountScopes, 0), Conjunction(conjunction),
+        AfterConjunction(erase(conjunction)), OuterSolutions(solutions) {
     assert(conjunction->getKind() == ConstraintKind::Conjunction);
 
     // Make a snapshot of the constraint system state before conjunction.
     if (conjunction->isIsolated())
       Snapshot.emplace(cs, conjunction);
+
+    if (cs.Timer) {
+      auto remainingTime = cs.Timer->getRemainingProcessTimeInMillis();
+      OuterTimeRemaining.emplace(cs.Timer->getAnchor(), remainingTime);
+    }
   }
 
   ~ConjunctionStep() override {
@@ -911,6 +925,12 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
     // successful outcome should keep a score set by `restoreOuterState`.
     if (HadFailure)
       restoreOriginalScores();
+
+    if (OuterTimeRemaining) {
+      auto anchor = OuterTimeRemaining->first;
+      auto remainingTime = OuterTimeRemaining->second;
+      CS.Timer.emplace(anchor, CS, remainingTime);
+    }
   }
 
   StepResult resume(bool prevFailed) override;

lib/Sema/ConstraintSystem.cpp

@@ -43,13 +43,18 @@ using namespace inference;
 
 #define DEBUG_TYPE "ConstraintSystem"
 
-ExpressionTimer::ExpressionTimer(Expr *E, ConstraintSystem &CS)
-  : E(E),
-      Context(CS.getASTContext()),
+ExpressionTimer::ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS)
+    : ExpressionTimer(
+          Anchor, CS,
+          CS.getASTContext().TypeCheckerOpts.ExpressionTimeoutThreshold) {}
+
+ExpressionTimer::ExpressionTimer(AnchorType Anchor, ConstraintSystem &CS,
+                                 unsigned thresholdInMillis)
+    : Anchor(Anchor), Context(CS.getASTContext()),
       StartTime(llvm::TimeRecord::getCurrentTime()),
+      ThresholdInMillis(thresholdInMillis),
       PrintDebugTiming(CS.getASTContext().TypeCheckerOpts.DebugTimeExpressions),
-      PrintWarning(true) {
-}
+      PrintWarning(true) {}
 
 ExpressionTimer::~ExpressionTimer() {
   auto elapsed = getElapsedProcessTimeInFractionalSeconds();
@@ -59,21 +64,39 @@ ExpressionTimer::~ExpressionTimer() {
     // Round up to the nearest 100th of a millisecond.
     llvm::errs() << llvm::format("%0.2f", ceil(elapsed * 100000) / 100)
                  << "ms\t";
-    E->getLoc().print(llvm::errs(), Context.SourceMgr);
+    if (auto *E = Anchor.dyn_cast<Expr *>()) {
+      E->getLoc().print(llvm::errs(), Context.SourceMgr);
+    } else {
+      auto *locator = Anchor.get<ConstraintLocator *>();
+      locator->dump(&Context.SourceMgr, llvm::errs());
+    }
     llvm::errs() << "\n";
   }
 
   if (!PrintWarning)
     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 *>();
+  }
+
   const auto WarnLimit = getWarnLimit();
-  if (WarnLimit != 0 && elapsedMS >= WarnLimit && E->getLoc().isValid())
-    Context.Diags.diagnose(E->getLoc(), diag::debug_long_expression,
-                           elapsedMS, WarnLimit)
-      .highlight(E->getSourceRange());
+  if (WarnLimit != 0 && elapsedMS >= WarnLimit &&
+      anchor.getStartLoc().isValid()) {
+    Context.Diags
+        .diagnose(anchor.getStartLoc(), diag::debug_long_expression, elapsedMS,
+                  WarnLimit)
+        .highlight(anchor.getSourceRange());
+  }
 }
 
-
 ConstraintSystem::ConstraintSystem(DeclContext *dc,
                                    ConstraintSystemOptions options)
   : Context(dc->getASTContext()), DC(dc), Options(options),