Sema: Record favored constraints in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -84,6 +84,8 @@ class SolverTrail {
     RecordedKeyPathComponentType,
     /// Disabled a constraint.
     DisabledConstraint,
+    /// Favored a constraint.
+    FavoredConstraint,
   };
 
   /// A change made to the constraint system.
@@ -248,6 +250,9 @@ class SolverTrail {
     /// Create a change that disabled a constraint.
     static Change disabledConstraint(Constraint *constraint);
 
+    /// Create a change that favored a constraint.
+    static Change favoredConstraint(Constraint *constraint);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2631,15 +2631,6 @@ class ConstraintSystem {
       }
 
       generatedConstraints.erase(genStart, genEnd);
-
-      for (unsigned constraintIdx :
-             range(scope->numFavoredConstraints, favoredConstraints.size())) {
-        if (favoredConstraints[constraintIdx]->isFavored())
-          favoredConstraints[constraintIdx]->setFavored(false);
-      }
-      favoredConstraints.erase(
-          favoredConstraints.begin() + scope->numFavoredConstraints,
-          favoredConstraints.end());
     }
 
     /// Check whether constraint system is allowed to form solutions
@@ -2651,21 +2642,18 @@ class ConstraintSystem {
     /// Disable the given constraint; this change will be rolled back
     /// when we exit the current solver scope.
     void disableConstraint(Constraint *constraint) {
+      ASSERT(!constraint->isDisabled());
       constraint->setDisabled();
       Trail.recordChange(SolverTrail::Change::disabledConstraint(constraint));
     }
 
-    unsigned getNumFavoredConstraints() const {
-      return favoredConstraints.size();
-    }
-
     /// Favor the given constraint; this change will be rolled back
     /// when we exit the current solver scope.
     void favorConstraint(Constraint *constraint) {
       assert(!constraint->isFavored());
 
       constraint->setFavored();
-      favoredConstraints.push_back(constraint);
+      Trail.recordChange(SolverTrail::Change::favoredConstraint(constraint));
     }
 
   private:
@@ -2689,8 +2677,7 @@ class ConstraintSystem {
     llvm::SmallVector<
       std::tuple<SolverScope *, ConstraintList::iterator, unsigned>, 4> scopes;
 
-    SmallVector<Constraint *, 4> favoredConstraints;
-    
+
     /// Depth of the solution stack.
     unsigned depth = 0;
   };
@@ -2871,8 +2858,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    unsigned numFavoredConstraints;
-
     unsigned numResultBuilderTransformed;
 
     /// The length of \c appliedPropertyWrappers

lib/Sema/CSSolver.cpp

@@ -664,7 +664,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numFavoredConstraints = cs.solverState->getNumFavoredConstraints();
   numResultBuilderTransformed = cs.resultBuilderTransformed.size();
   numAppliedPropertyWrappers = cs.appliedPropertyWrappers.size();
   numResolvedOverloads = cs.ResolvedOverloads.size();

lib/Sema/CSTrail.cpp

@@ -254,6 +254,14 @@ SolverTrail::Change::disabledConstraint(Constraint *constraint) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::favoredConstraint(Constraint *constraint) {
+  Change result;
+  result.Kind = ChangeKind::FavoredConstraint;
+  result.TheConstraint.Constraint = constraint;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -355,6 +363,11 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     if (TheConstraint.Constraint->isDisabled())
       TheConstraint.Constraint->setEnabled();
     break;
+
+  case ChangeKind::FavoredConstraint:
+    if (TheConstraint.Constraint->isFavored())
+      TheConstraint.Constraint->setFavored(false);
+    break;
   }
 }
 
@@ -550,6 +563,13 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
                                     indent + 2);
     out << ")\n";
     break;
+
+  case ChangeKind::FavoredConstraint:
+    out << "(favored constraint ";
+    TheConstraint.Constraint->print(out, &cs.getASTContext().SourceMgr,
+                                    indent + 2);
+    out << ")\n";
+    break;
   }
 }