Sema: Record score increases in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -78,8 +78,10 @@ CHANGE(RecordedCaseLabelItemInfo)
 CHANGE(RecordedPotentialThrowSite)
 CHANGE(RecordedIsolatedParam)
 CHANGE(RecordedKeyPath)
+CHANGE(IncreasedScore)
+CHANGE(DecreasedScore)
 
-LAST_CHANGE(RecordedKeyPath)
+LAST_CHANGE(DecreasedScore)
 
 #undef LOCATOR_CHANGE
 #undef EXPR_CHANGE

include/swift/Sema/CSTrail.h

@@ -229,6 +229,12 @@ class SolverTrail {
     /// Create a change that recorded a key path expression.
     static Change RecordedKeyPath(KeyPathExpr *expr);
 
+    /// Create a change that increased the score.
+    static Change IncreasedScore(ScoreKind kind, unsigned value);
+
+    /// Create a change that decreased the score.
+    static Change DecreasedScore(ScoreKind kind, unsigned value);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2857,9 +2857,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The previous score.
-    Score PreviousScore;
-
     /// The scope number of this scope. Set when the scope is registered.
     unsigned scopeNumber = 0;
 
@@ -5534,10 +5531,13 @@ class ConstraintSystem {
 
 public:
   /// Increase the score of the given kind for the current (partial) solution
-  /// along the.
+  /// along the current solver path.
   void increaseScore(ScoreKind kind, ConstraintLocatorBuilder Locator,
                      unsigned value = 1);
 
+  /// Primitive form of the above. Records a change in the trail.
+  void increaseScore(ScoreKind kind, unsigned value);
+
   /// Determine whether this solution is guaranteed to be worse than the best
   /// solution found so far.
   bool worseThanBestSolution() const;

lib/Sema/CSRanking.cpp

@@ -110,6 +110,14 @@ static bool shouldIgnoreScoreIncreaseForCodeCompletion(
   return false;
 }
 
+void ConstraintSystem::increaseScore(ScoreKind kind, unsigned value) {
+  unsigned index = static_cast<unsigned>(kind);
+  CurrentScore.Data[index] += value;
+
+  if (solverState && value > 0)
+    recordChange(SolverTrail::Change::IncreasedScore(kind, value));
+}
+
 void ConstraintSystem::increaseScore(ScoreKind kind,
                                      ConstraintLocatorBuilder Locator,
                                      unsigned value) {
@@ -135,8 +143,7 @@ void ConstraintSystem::increaseScore(ScoreKind kind,
     llvm::errs() << ")\n";
   }
 
-  unsigned index = static_cast<unsigned>(kind);
-  CurrentScore.Data[index] += value;
+  increaseScore(kind, value);
 }
 
 bool ConstraintSystem::worseThanBestSolution() const {

lib/Sema/CSSolver.cpp

@@ -274,8 +274,15 @@ Solution ConstraintSystem::finalize() {
 
 void ConstraintSystem::replaySolution(const Solution &solution,
                                       bool shouldIncreaseScore) {
-  if (shouldIncreaseScore)
-    CurrentScore += solution.getFixedScore();
+  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);
+    }
+  }
 
   // Assign fixed types to the type variables solved by this solution.
   for (auto binding : solution.typeBindings) {
@@ -711,8 +718,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
 
-  PreviousScore = cs.CurrentScore;
-
   cs.solverState->registerScope(this);
   assert(!cs.failedConstraint && "Unexpected failed constraint!");
 }
@@ -742,9 +747,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  // Reset the previous score.
-  cs.CurrentScore = PreviousScore;
-
   // Clear out other "failed" state.
   cs.failedConstraint = nullptr;
 }

lib/Sema/CSStep.cpp

@@ -885,6 +885,12 @@ 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();
 
   // Reset the scope counter to avoid "too complex" failures

lib/Sema/CSStep.h

@@ -906,8 +906,6 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
 
   /// Best solution solver reached so far.
   std::optional<Score> BestScore;
-  /// The score established before conjunction is attempted.
-  Score CurrentScore;
 
   /// The number of constraint solver scopes already explored
   /// before accepting this conjunction.
@@ -949,7 +947,7 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
                   SmallVectorImpl<Solution> &solutions)
       : BindingStep(cs, {cs, conjunction},
                     conjunction->isIsolated() ? IsolatedSolutions : solutions),
-        BestScore(getBestScore()), CurrentScore(getCurrentScore()),
+        BestScore(getBestScore()),
         OuterScopeCount(cs.CountScopes, 0), Conjunction(conjunction),
         AfterConjunction(erase(conjunction)), OuterSolutions(solutions) {
     assert(conjunction->getKind() == ConstraintKind::Conjunction);
@@ -975,11 +973,8 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
 
     // Restore best score only if conjunction fails because
     // successful outcome should keep a score set by `restoreOuterState`.
-    if (HadFailure) {
-      auto solutionScore = Score();
+    if (HadFailure)
       restoreBestScore();
-      restoreCurrentScore(solutionScore);
-    }
 
     if (OuterTimeRemaining) {
       auto anchor = OuterTimeRemaining->first;
@@ -1033,7 +1028,6 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
 private:
   /// Restore best and current scores as they were before conjunction.
   void restoreCurrentScore(const Score &solutionScore) const {
-    CS.CurrentScore = CurrentScore;
     CS.increaseScore(SK_Fix, Conjunction->getLocator(),
                      solutionScore.Data[SK_Fix]);
     CS.increaseScore(SK_Hole, Conjunction->getLocator(),

lib/Sema/CSTrail.cpp

@@ -325,6 +325,26 @@ SolverTrail::Change::RecordedKeyPath(KeyPathExpr *expr) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::IncreasedScore(ScoreKind kind, unsigned value) {
+  ASSERT(value <= 0xffffff && "value must fit in 24 bits");
+
+  Change result;
+  result.Kind = ChangeKind::IncreasedScore;
+  result.Options = unsigned(kind) | (value << 8);
+  return result;
+}
+
+SolverTrail::Change
+SolverTrail::Change::DecreasedScore(ScoreKind kind, unsigned value) {
+  ASSERT(value <= 0xffffff && "value must fit in 24 bits");
+
+  Change result;
+  result.Kind = ChangeKind::DecreasedScore;
+  result.Options = unsigned(kind) | (value << 8);
+  return result;
+}
+
 SyntacticElementTargetKey
 SolverTrail::Change::getSyntacticElementTargetKey() const {
   ASSERT(Kind == ChangeKind::RecordedTarget);
@@ -487,6 +507,21 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedKeyPath:
     cs.removeKeyPath(KeyPath.Expr);
     break;
+
+  case ChangeKind::IncreasedScore: {
+    auto kind = Options & 0xff;
+    unsigned value = Options >> 8;
+    ASSERT(cs.CurrentScore.Data[kind] >= value);
+    cs.CurrentScore.Data[kind] -= value;
+    break;
+  }
+
+  case ChangeKind::DecreasedScore: {
+    auto kind = Options & 0xff;
+    unsigned value = Options >> 8;
+    cs.CurrentScore.Data[kind] += value;
+    break;
+  }
   }
 }
 
@@ -706,6 +741,18 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     simple_display(out, KeyPath.Expr);
     out << ")\n";
     break;
+
+  case ChangeKind::IncreasedScore:
+    out << "(IncreasedScore ";
+    out << Score::getNameFor(ScoreKind(Options & 0xff));
+    out << " by " << (Options >> 8) << ")\n";
+    break;
+
+  case ChangeKind::DecreasedScore:
+    out << "(DecreasedScore ";
+    out << Score::getNameFor(ScoreKind(Options & 0xff));
+    out << " by " << (Options >> 8) << ")\n";
+    break;
   }
 }