Merge pull request #76932 from slavapestov/cstrail-part-6

slavapestov · web-flow · commit 04f48db7f207 · 2024-10-09T20:39:09.000-04:00
Sema: Small SolverTrail cleanups
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -2143,8 +2143,6 @@ class ConstraintSystem {
   friend struct ClosureIsolatedByPreconcurrency;
   friend class SolverTrail;
 
-  class SolverScope;
-
   /// Expressions that are known to be unevaluated.
   /// Note: this is only used to support ObjCSelectorExpr at the moment.
   llvm::SmallPtrSet<Expr *, 2> UnevaluatedRootExprs;
@@ -2555,23 +2553,21 @@ class ConstraintSystem {
     }
 
     /// Update statistics when a scope begins.
-    ///
-    /// \param scope The scope to associate with current solver state.
-    void beginScope(SolverScope *scope) {
+    unsigned beginScope() {
       ++depth;
       maxDepth = std::max(maxDepth, depth);
-      scope->scopeNumber = NumStatesExplored++;
 
       CS.incrementScopeCounter();
+
+      return NumStatesExplored++;
     }
 
     /// Update statistics when a scope ends.
-    ///
-    /// \param scope The solver scope to rollback.
-    void endScope(SolverScope *scope) {
+    void endScope(unsigned scopeNumber) {
+      ASSERT(depth > 0);
       --depth;
 
-      unsigned countScopesExplored = NumStatesExplored - scope->scopeNumber;
+      unsigned countScopesExplored = NumStatesExplored - scopeNumber;
       if (countScopesExplored == 1)
         CS.incrementLeafScopes();
     }
@@ -2775,7 +2771,7 @@ class ConstraintSystem {
   /// this object is destroyed.
   ///
   ///
-  class SolverScope {
+  struct SolverScope {
     ConstraintSystem &cs;
 
     /// The length of \c TypeVariables.
@@ -2785,15 +2781,19 @@ class ConstraintSystem {
     unsigned numTrailChanges;
 
     /// The scope number of this scope. Set when the scope is registered.
-    unsigned scopeNumber = 0;
+    unsigned scopeNumber : 31;
+
+    /// A moved-from scope skips doing anything in the destructor.
+    unsigned moved : 1;
+
+    explicit SolverScope(ConstraintSystem &cs);
 
     SolverScope(const SolverScope &) = delete;
-    SolverScope &operator=(const SolverScope &) = delete;
+    SolverScope(SolverScope &&other);
 
-    friend class ConstraintSystem;
+    SolverScope &operator=(const SolverScope &) = delete;
+    SolverScope &operator=(SolverScope &&) = delete;
 
-  public:
-    explicit SolverScope(ConstraintSystem &cs);
     ~SolverScope();
   };
 
@@ -5465,6 +5465,15 @@ class ConstraintSystem {
   /// Primitive form of the above. Records a change in the trail.
   void increaseScore(ScoreKind kind, unsigned value);
 
+  /// Apply the score from a partial solution. Records the change in the
+  /// trail.
+  void replayScore(const Score &score);
+
+  /// Temporarily zero out the score, and record this in the trail so that
+  /// we restore the score when the scope ends. Used when solving a
+  /// ConjunctionStep.
+  void clearScore();
+
   /// Determine whether this solution is guaranteed to be worse than the best
   /// solution found so far.
   bool worseThanBestSolution() const;
diff --git a/lib/Sema/CSRanking.cpp b/lib/Sema/CSRanking.cpp
@@ -146,6 +146,27 @@ void ConstraintSystem::increaseScore(ScoreKind kind,
   increaseScore(kind, value);
 }
 
+void ConstraintSystem::replayScore(const Score &score) {
+  if (solverState) {
+    for (unsigned i = 0; i < NumScoreKinds; ++i) {
+      if (unsigned value = score.Data[i])
+        recordChange(
+          SolverTrail::Change::IncreasedScore(ScoreKind(i), value));
+    }
+  }
+  CurrentScore += score;
+}
+
+void ConstraintSystem::clearScore() {
+  for (unsigned i = 0; i < NumScoreKinds; ++i) {
+    if (unsigned value = CurrentScore.Data[i]) {
+      recordChange(
+        SolverTrail::Change::DecreasedScore(ScoreKind(i), value));
+    }
+  }
+  CurrentScore = Score();
+}
+
 bool ConstraintSystem::worseThanBestSolution() const {
   if (getASTContext().TypeCheckerOpts.DisableConstraintSolverPerformanceHacks)
     return false;
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -271,15 +271,8 @@ Solution ConstraintSystem::finalize() {
 
 void ConstraintSystem::replaySolution(const Solution &solution,
                                       bool shouldIncreaseScore) {
-  if (shouldIncreaseScore) {
-    // Update the score. We do this instead of operator+= because we
-    // want to record the increments in the trail.
-    auto solutionScore = solution.getFixedScore();
-    for (unsigned i = 0; i < NumScoreKinds; ++i) {
-      if (unsigned value = solutionScore.Data[i])
-        increaseScore(ScoreKind(i), value);
-    }
-  }
+  if (shouldIncreaseScore)
+    replayScore(solution.getFixedScore());
 
   // Assign fixed types to the type variables solved by this solution.
   for (auto binding : solution.typeBindings) {
@@ -702,21 +695,32 @@ ConstraintSystem::SolverState::~SolverState() {
 }
 
 ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
-  : cs(cs) {
-  numTrailChanges = cs.solverState->Trail.size();
-
-  numTypeVariables = cs.TypeVariables.size();
+  : cs(cs),
+    numTypeVariables(cs.TypeVariables.size()),
+    numTrailChanges(cs.solverState->Trail.size()),
+    scopeNumber(cs.solverState->beginScope()),
+    moved(0) {
+  ASSERT(!cs.failedConstraint && "Unexpected failed constraint!");
+}
 
-  cs.solverState->beginScope(this);
-  assert(!cs.failedConstraint && "Unexpected failed constraint!");
+ConstraintSystem::SolverScope::SolverScope(SolverScope &&other)
+  : cs(other.cs),
+    numTypeVariables(other.numTypeVariables),
+    numTrailChanges(other.numTrailChanges),
+    scopeNumber(other.scopeNumber),
+    moved(0) {
+  other.moved = 1;
 }
 
 ConstraintSystem::SolverScope::~SolverScope() {
+  if (moved)
+    return;
+
   // Don't attempt to rollback from an incorrect state.
   if (cs.inInvalidState())
     return;
 
-  // Erase the end of various lists.
+  // Roll back introduced type variables.
   truncate(cs.TypeVariables, numTypeVariables);
 
   // Move any remaining active constraints into the inactive list.
@@ -731,7 +735,8 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Roll back changes to the constraint system.
   cs.solverState->Trail.undo(numTrailChanges);
 
-  cs.solverState->endScope(this);
+  // Update statistics.
+  cs.solverState->endScope(scopeNumber);
 
   // Clear out other "failed" state.
   cs.failedConstraint = nullptr;
diff --git a/lib/Sema/CSStep.cpp b/lib/Sema/CSStep.cpp
@@ -885,13 +885,7 @@ bool ConjunctionStep::attempt(const ConjunctionElement &element) {
 
   // Make sure that element is solved in isolation
   // by dropping all scoring information.
-  for (unsigned i = 0; i < NumScoreKinds; ++i) {
-    if (unsigned value = CS.CurrentScore.Data[i]) {
-      CS.recordChange(
-        SolverTrail::Change::DecreasedScore(ScoreKind(i), value));
-    }
-  }
-  CS.CurrentScore = Score();
+  CS.clearScore();
 
   // Reset the scope counter to avoid "too complex" failures
   // when closure has a lot of elements in the body.
@@ -1039,7 +1033,7 @@ StepResult ConjunctionStep::resume(bool prevFailed) {
           // restored right afterwards because score of the
           // element does contribute to the overall score.
           restoreBestScore();
-          restoreCurrentScore(solution.getFixedScore());
+          updateScoreAfterConjunction(solution.getFixedScore());
 
           // Transform all of the unbound outer variables into
           // placeholders since we are not going to solve for
@@ -1091,10 +1085,10 @@ StepResult ConjunctionStep::resume(bool prevFailed) {
 }
 
 void ConjunctionStep::restoreOuterState(const Score &solutionScore) const {
-  // Restore best/current score, since upcoming step is going to
-  // work with outer scope in relation to the conjunction.
+  // Restore best score and update current score, since upcoming step
+  // is going to work with outer scope in relation to the conjunction.
   restoreBestScore();
-  restoreCurrentScore(solutionScore);
+  updateScoreAfterConjunction(solutionScore);
 
   // Active all of the previously out-of-scope constraints
   // because conjunction can propagate type information up
diff --git a/lib/Sema/CSStep.h b/lib/Sema/CSStep.h
@@ -519,7 +519,7 @@ template <typename P> class BindingStep : public SolverStep {
   /// being attempted, helps to rewind state of the
   /// constraint system back to original before attempting
   /// next binding, if any.
-  std::optional<std::pair<std::unique_ptr<Scope>, typename P::Element>>
+  std::optional<std::pair<Scope, typename P::Element>>
       ActiveChoice;
 
   BindingStep(ConstraintSystem &cs, P producer,
@@ -548,16 +548,14 @@ template <typename P> class BindingStep : public SolverStep {
       }
 
       {
-        auto &trail = CS.solverState->Trail;
-        unsigned size = trail.size();
-
-        auto scope = std::make_unique<Scope>(CS);
+        Scope scope(CS);
         if (attempt(*choice)) {
           ActiveChoice.emplace(std::move(scope), *choice);
 
           if (CS.isDebugMode()) {
-            trail.dumpActiveScopeChanges(
-              llvm::errs(), size, CS.solverState->getCurrentIndent());
+            CS.solverState->Trail.dumpActiveScopeChanges(
+              llvm::errs(), ActiveChoice->first.numTrailChanges,
+              CS.solverState->getCurrentIndent());
           }
           
           return suspend(std::make_unique<SplitterStep>(CS, Solutions));
@@ -1029,8 +1027,10 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
   }
 
 private:
-  /// Restore best and current scores as they were before conjunction.
-  void restoreCurrentScore(const Score &solutionScore) const {
+  /// We need to do this to make sure that we rank solutions with
+  /// invalid closures appropriately and don’t produce a valid
+  /// solution if a multi-statement closure failed.
+  void updateScoreAfterConjunction(const Score &solutionScore) const {
     CS.increaseScore(SK_Fix, Conjunction->getLocator(),
                      solutionScore.Data[SK_Fix]);
     CS.increaseScore(SK_Hole, Conjunction->getLocator(),
diff --git a/lib/Sema/ConstraintGraph.cpp b/lib/Sema/ConstraintGraph.cpp
@@ -1282,7 +1282,7 @@ bool ConstraintGraph::contractEdges() {
         continue;
 
       // This closure is not currently in scope.
-      if (!CS.TypeVariables.count(paramTy))
+      if (!CS.isActiveTypeVariable(paramTy))
         break;
 
       // Nothing to contract here since outside parameter
