Sema: Record the trail step count in solver statistics

Also introduce two new frontend flags:

The -solver-scope-threshold flag sets the maximum number of scopes, which was
previously hardcoded to 1 million.

The -solver-trail-threshold flag sets the maximum number of trail steps,
which defaults to 64 million.

Author: slavapestov

include/swift/Basic/LangOptions.h

@@ -847,6 +847,16 @@ namespace swift {
     /// than this many seconds.
     unsigned ExpressionTimeoutThreshold = 600;
 
+    /// The upper bound, in bytes, of temporary data that can be
+    /// allocated by the constraint solver.
+    unsigned SolverMemoryThreshold = 512 * 1024 * 1024;
+
+    /// The maximum number of scopes we explore before giving up.
+    unsigned SolverScopeThreshold = 1024 * 1024;
+
+    /// The maximum number of trail steps we take before giving up.
+    unsigned SolverTrailThreshold = 64 * 1024 * 1024;
+
     /// If non-zero, abort the switch statement exhaustiveness checker if
     /// the Space::minus function is called more than this many times.
     ///
@@ -896,12 +906,6 @@ namespace swift {
     /// is for testing purposes.
     std::vector<std::string> DebugForbidTypecheckPrefixes;
 
-    /// The upper bound, in bytes, of temporary data that can be
-    /// allocated by the constraint solver.
-    unsigned SolverMemoryThreshold = 512 * 1024 * 1024;
-
-    unsigned SolverBindingThreshold = 1024 * 1024;
-
     /// The upper bound to number of sub-expressions unsolved
     /// before termination of the shrink phrase of the constraint solver.
     unsigned SolverShrinkUnsolvedThreshold = 10;

include/swift/Basic/Statistics.def

@@ -184,9 +184,14 @@ FRONTEND_STATISTIC(Sema, NumCrossImportsChecked)
 /// Number of pairs of modules we've actually found cross-imports for.
 FRONTEND_STATISTIC(Sema, NumCrossImportsFound)
 
+/// Number of steps recorded in the trail while solving expression type
+/// constraints. A rough proxy for "how much work the expression
+/// type checker did".
+FRONTEND_STATISTIC(Sema, NumTrailSteps)
+
 /// Number of constraint-solving scopes created in the typechecker, while
-/// solving expression type constraints. A rough proxy for "how much work the
-/// expression typechecker did".
+/// solving expression type constraints. Another rough proxy for "how much
+/// work the expression type checker did".
 FRONTEND_STATISTIC(Sema, NumConstraintScopes)
 
 /// Number of constraint-solving scopes that were created but which

include/swift/Option/FrontendOptions.td

@@ -827,6 +827,10 @@ def enable_volatile_modules : Flag<["-"], "enable-volatile-modules">,
 
 def solver_expression_time_threshold_EQ : Joined<["-"], "solver-expression-time-threshold=">;
 
+def solver_scope_threshold_EQ : Joined<["-"], "solver-scope-threshold=">;
+
+def solver_trail_threshold_EQ : Joined<["-"], "solver-trail-threshold=">;
+
 def solver_disable_shrink :
   Flag<["-"], "solver-disable-shrink">,
   HelpText<"Disable the shrink phase of expression type checking">;

include/swift/Sema/ConstraintSolverStats.def

@@ -26,6 +26,7 @@ CS_STATISTIC(NumConjunctionTerms, "# of conjunction terms explored")
 CS_STATISTIC(NumSimplifiedConstraints, "# of constraints simplified")
 CS_STATISTIC(NumUnsimplifiedConstraints, "# of constraints not simplified")
 CS_STATISTIC(NumSimplifyIterations, "# of simplification iterations")
-CS_STATISTIC(NumStatesExplored, "# of solution states explored")
+CS_STATISTIC(NumSolverScopes, "# of solution states explored")
+CS_STATISTIC(NumTrailSteps, "# of changes recorded in the trail")
 CS_STATISTIC(NumComponentsSplit, "# of connected components split")
 #undef CS_STATISTIC

include/swift/Sema/ConstraintSystem.h

@@ -2170,13 +2170,20 @@ class ConstraintSystem {
   /// Counter for type variables introduced.
   unsigned TypeCounter = 0;
 
+  /// The number of changes recorded in the trail so far during the
+  /// solution of this constraint system.
+  ///
+  /// This is a rough proxy for how much work the solver did.
+  unsigned NumTrailSteps = 0;
+
   /// The number of scopes created so far during the solution
   /// of this constraint system.
   ///
-  /// This is a measure of complexity of the solution space. A new
-  /// scope is created every time we attempt a type variable binding
-  /// or explore an option in a disjunction.
-  unsigned CountScopes = 0;
+  /// A new scope is created every time we attempt a type variable
+  /// binding or explore an option in a disjunction.
+  ///
+  /// This is a measure of complexity of the solution space.
+  unsigned NumSolverScopes = 0;
 
   /// High-water mark of measured memory usage in any sub-scope we
   /// explored.
@@ -2539,24 +2546,12 @@ class ConstraintSystem {
     }
 
     /// Update statistics when a scope begins.
-    unsigned beginScope() {
-      ++depth;
-      maxDepth = std::max(maxDepth, depth);
-
-      CS.incrementScopeCounter();
-
-      return NumStatesExplored++;
-    }
+    unsigned beginScope();
 
     /// Update statistics when a scope ends.
-    void endScope(unsigned scopeNumber) {
-      ASSERT(depth > 0);
-      --depth;
-
-      unsigned countScopesExplored = NumStatesExplored - scopeNumber;
-      if (countScopesExplored == 1)
-        CS.incrementLeafScopes();
-    }
+    void endScope(unsigned scopeNumber,
+                  uint64_t startTrailSteps,
+                  uint64_t endTrailSteps);
 
     /// Check whether constraint system is allowed to form solutions
     /// even with unbound type variables present.
@@ -2760,11 +2755,11 @@ class ConstraintSystem {
   struct SolverScope {
     ConstraintSystem &cs;
 
-    /// The length of \c TypeVariables.
-    unsigned numTypeVariables;
+    /// The length of \c TypeVariables at the start of the scope.
+    unsigned startTypeVariables;
 
-    /// The length of \c Trail.
-    unsigned numTrailChanges;
+    /// The length of \c Trail at the start of the scope.
+    uint64_t startTrailSteps;
 
     /// The scope number of this scope. Set when the scope is registered.
     unsigned scopeNumber : 31;
@@ -5538,9 +5533,16 @@ class ConstraintSystem {
       return true;
     }
 
-    // Bail out once we've looked at a really large number of
-    // choices.
-    if (CountScopes > getASTContext().TypeCheckerOpts.SolverBindingThreshold) {
+    auto &opts = getASTContext().TypeCheckerOpts;
+
+    // Bail out once we've looked at a really large number of choices.
+    if (NumSolverScopes > opts.SolverScopeThreshold) {
+      isAlreadyTooComplex.first = true;
+      return true;
+    }
+
+    // Bail out once we've taken a really large number of steps.
+    if (NumTrailSteps > opts.SolverTrailThreshold) {
       isAlreadyTooComplex.first = true;
       return true;
     }

lib/Frontend/CompilerInvocation.cpp

@@ -1777,6 +1777,10 @@ static bool ParseTypeCheckerArgs(TypeCheckerOptions &Opts, ArgList &Args,
                              Opts.DebugConstraintSolverAttempt);
   setUnsignedIntegerArgument(OPT_solver_memory_threshold,
                              Opts.SolverMemoryThreshold);
+  setUnsignedIntegerArgument(OPT_solver_scope_threshold_EQ,
+                             Opts.SolverScopeThreshold);
+  setUnsignedIntegerArgument(OPT_solver_trail_threshold_EQ,
+                             Opts.SolverTrailThreshold);
   setUnsignedIntegerArgument(OPT_solver_shrink_unsolved_threshold,
                              Opts.SolverShrinkUnsolvedThreshold);
 

lib/Sema/CSSolver.cpp

@@ -682,7 +682,7 @@ ConstraintSystem::SolverState::~SolverState() {
   // Update the "largest" statistics if this system is larger than the
   // previous one.  
   // FIXME: This is not at all thread-safe.
-  if (NumStatesExplored > LargestNumStatesExplored.getValue()) {
+  if (NumSolverScopes > LargestNumSolverScopes.getValue()) {
     LargestSolutionAttemptNumber = SolutionAttempt-1;
     ++LargestSolutionAttemptNumber;
     #define CS_STATISTIC(Name, Description) \
@@ -695,17 +695,17 @@ ConstraintSystem::SolverState::~SolverState() {
 
 ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   : cs(cs),
-    numTypeVariables(cs.TypeVariables.size()),
-    numTrailChanges(cs.solverState->Trail.size()),
+    startTypeVariables(cs.TypeVariables.size()),
+    startTrailSteps(cs.solverState->Trail.size()),
     scopeNumber(cs.solverState->beginScope()),
     moved(0) {
   ASSERT(!cs.failedConstraint && "Unexpected failed constraint!");
 }
 
 ConstraintSystem::SolverScope::SolverScope(SolverScope &&other)
   : cs(other.cs),
-    numTypeVariables(other.numTypeVariables),
-    numTrailChanges(other.numTrailChanges),
+    startTypeVariables(other.startTypeVariables),
+    startTrailSteps(other.startTrailSteps),
     scopeNumber(other.scopeNumber),
     moved(0) {
   other.moved = 1;
@@ -720,7 +720,7 @@ ConstraintSystem::SolverScope::~SolverScope() {
     return;
 
   // Roll back introduced type variables.
-  truncate(cs.TypeVariables, numTypeVariables);
+  truncate(cs.TypeVariables, startTypeVariables);
 
   // Move any remaining active constraints into the inactive list.
   if (!cs.ActiveConstraints.empty()) {
@@ -731,16 +731,43 @@ ConstraintSystem::SolverScope::~SolverScope() {
                                   cs.ActiveConstraints);
   }
 
+  uint64_t endTrailSteps = cs.solverState->Trail.size();
+
   // Roll back changes to the constraint system.
-  cs.solverState->Trail.undo(numTrailChanges);
+  cs.solverState->Trail.undo(startTrailSteps);
 
   // Update statistics.
-  cs.solverState->endScope(scopeNumber);
+  cs.solverState->endScope(scopeNumber,
+                           startTrailSteps,
+                           endTrailSteps);
 
   // Clear out other "failed" state.
   cs.failedConstraint = nullptr;
 }
 
+unsigned ConstraintSystem::SolverState::beginScope() {
+  ++depth;
+  maxDepth = std::max(maxDepth, depth);
+
+  CS.incrementScopeCounter();
+
+  return NumSolverScopes++;
+}
+
+/// Update statistics when a scope ends.
+void ConstraintSystem::SolverState::endScope(unsigned scopeNumber,
+                                             uint64_t startTrailSteps,
+                                             uint64_t endTrailSteps) {
+  ASSERT(depth > 0);
+  --depth;
+
+  NumTrailSteps += (endTrailSteps - startTrailSteps);
+
+  unsigned countSolverScopes = NumSolverScopes - scopeNumber;
+  if (countSolverScopes == 1)
+    CS.incrementLeafScopes();
+}
+
 /// Solve the system of constraints.
 ///
 /// \param allowFreeTypeVariables How to bind free type variables in
@@ -1529,7 +1556,8 @@ bool ConstraintSystem::solve(SmallVectorImpl<Solution> &solutions,
   if (isDebugMode()) {
     auto &log = llvm::errs();
     log << "\n---Solver statistics---\n";
-    log << "Total number of scopes explored: " << solverState->NumStatesExplored << "\n";
+    log << "Total number of scopes explored: " << solverState->NumSolverScopes << "\n";
+    log << "Total number of trail steps: " << solverState->NumTrailSteps << "\n";
     log << "Maximum depth reached while exploring solutions: " << solverState->maxDepth << "\n";
     if (Timer) {
       auto timeInMillis =

lib/Sema/CSStep.cpp

@@ -891,9 +891,10 @@ bool ConjunctionStep::attempt(const ConjunctionElement &element) {
   // by dropping all scoring information.
   CS.clearScore();
 
-  // Reset the scope counter to avoid "too complex" failures
-  // when closure has a lot of elements in the body.
-  CS.CountScopes = 0;
+  // Reset the scope and trail counters to avoid "too complex"
+  // failures when closure has a lot of elements in the body.
+  CS.NumSolverScopes = 0;
+  CS.NumTrailSteps = 0;
 
   // If timer is enabled, let's reset it so that each element
   // (expression) gets a fresh time slice to get solved. This

lib/Sema/CSStep.h

@@ -542,7 +542,7 @@ template <typename P> class BindingStep : public SolverStep {
 
           if (CS.isDebugMode()) {
             CS.solverState->Trail.dumpActiveScopeChanges(
-              llvm::errs(), ActiveChoice->first.numTrailChanges,
+              llvm::errs(), ActiveChoice->first.startTrailSteps,
               CS.solverState->getCurrentIndent());
           }
 
@@ -893,8 +893,12 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
   std::optional<Score> BestScore;
 
   /// The number of constraint solver scopes already explored
-  /// before accepting this conjunction.
-  llvm::SaveAndRestore<unsigned> OuterScopeCount;
+  /// before attempting this conjunction.
+  llvm::SaveAndRestore<unsigned> OuterNumSolverScopes;
+
+  /// The number of trail steps already recorded before attempting
+  /// this conjunction.
+  llvm::SaveAndRestore<unsigned> OuterNumTrailSteps;
 
   /// The number of milliseconds until outer constraint system
   /// is considered "too complex" if timer is enabled.
@@ -929,7 +933,9 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
       : BindingStep(cs, {cs, conjunction},
                     conjunction->isIsolated() ? IsolatedSolutions : solutions),
         BestScore(getBestScore()),
-        OuterScopeCount(cs.CountScopes, 0), Conjunction(conjunction),
+        OuterNumSolverScopes(cs.NumSolverScopes, 0),
+        OuterNumTrailSteps(cs.NumTrailSteps, 0),
+        Conjunction(conjunction),
         OuterSolutions(solutions) {
     assert(conjunction->getKind() == ConstraintKind::Conjunction);
 

lib/Sema/ConstraintSystem.cpp

@@ -140,8 +140,7 @@ ConstraintSystem::~ConstraintSystem() {
 }
 
 void ConstraintSystem::incrementScopeCounter() {
-  ++CountScopes;
-  // FIXME: (transitional) increment the redundant "always-on" counter.
+  ++NumSolverScopes;
   if (auto *Stats = getASTContext().Stats)
     ++Stats->getFrontendCounters().NumConstraintScopes;
 }