Sema: Record disabled constraints in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -82,6 +82,8 @@ class SolverTrail {
     RecordedNodeType,
     /// Recorded an assignment of a type to a keypath component.
     RecordedKeyPathComponentType,
+    /// Disabled a constraint.
+    DisabledConstraint,
   };
 
   /// A change made to the constraint system.
@@ -243,6 +245,9 @@ class SolverTrail {
                                                unsigned component,
                                                Type oldType);
 
+    /// Create a change that disabled a constraint.
+    static Change disabledConstraint(Constraint *constraint);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2632,15 +2632,6 @@ class ConstraintSystem {
 
       generatedConstraints.erase(genStart, genEnd);
 
-      for (unsigned constraintIdx :
-             range(scope->numDisabledConstraints, disabledConstraints.size())) {
-        if (disabledConstraints[constraintIdx]->isDisabled())
-          disabledConstraints[constraintIdx]->setEnabled();
-      }
-      disabledConstraints.erase(
-          disabledConstraints.begin() + scope->numDisabledConstraints,
-          disabledConstraints.end());
-
       for (unsigned constraintIdx :
              range(scope->numFavoredConstraints, favoredConstraints.size())) {
         if (favoredConstraints[constraintIdx]->isFavored())
@@ -2657,15 +2648,11 @@ class ConstraintSystem {
       return AllowFreeTypeVariables != FreeTypeVariableBinding::Disallow;
     }
 
-    unsigned getNumDisabledConstraints() const {
-      return disabledConstraints.size();
-    }
-
     /// Disable the given constraint; this change will be rolled back
     /// when we exit the current solver scope.
     void disableConstraint(Constraint *constraint) {
       constraint->setDisabled();
-      disabledConstraints.push_back(constraint);
+      Trail.recordChange(SolverTrail::Change::disabledConstraint(constraint));
     }
 
     unsigned getNumFavoredConstraints() const {
@@ -2702,7 +2689,6 @@ class ConstraintSystem {
     llvm::SmallVector<
       std::tuple<SolverScope *, ConstraintList::iterator, unsigned>, 4> scopes;
 
-    SmallVector<Constraint *, 4> disabledConstraints;
     SmallVector<Constraint *, 4> favoredConstraints;
 
     /// Depth of the solution stack.
@@ -2885,8 +2871,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    unsigned numDisabledConstraints;
-
     unsigned numFavoredConstraints;
 
     unsigned numResultBuilderTransformed;

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();
-  numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
   numFavoredConstraints = cs.solverState->getNumFavoredConstraints();
   numResultBuilderTransformed = cs.resultBuilderTransformed.size();
   numAppliedPropertyWrappers = cs.appliedPropertyWrappers.size();
@@ -1771,10 +1770,12 @@ ConstraintSystem::filterDisjunction(
     if (restoreOnFail)
       constraintsToRestoreOnFail.push_back(constraint);
 
-    if (solverState)
-      solverState->disableConstraint(constraint);
-    else
-      constraint->setDisabled();
+    if (!constraint->isDisabled()) {
+      if (solverState)
+        solverState->disableConstraint(constraint);
+      else
+        constraint->setDisabled();
+    }
   }
 
   switch (numEnabledTerms) {

lib/Sema/CSTrail.cpp

@@ -246,6 +246,14 @@ SolverTrail::Change::recordedKeyPathComponentType(const KeyPathExpr *expr,
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::disabledConstraint(Constraint *constraint) {
+  Change result;
+  result.Kind = ChangeKind::DisabledConstraint;
+  result.TheConstraint.Constraint = constraint;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -342,6 +350,11 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedKeyPathComponentType:
     cs.restoreType(KeyPath.Expr, Options, KeyPath.OldType);
     break;
+
+  case ChangeKind::DisabledConstraint:
+    if (TheConstraint.Constraint->isDisabled())
+      TheConstraint.Constraint->setEnabled();
+    break;
   }
 }
 
@@ -530,6 +543,13 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
       out << "null";
     out << " for component " << Options << ")\n";
     break;
+
+  case ChangeKind::DisabledConstraint:
+    out << "(disabled constraint ";
+    TheConstraint.Constraint->print(out, &cs.getASTContext().SourceMgr,
+                                    indent + 2);
+    out << ")\n";
+    break;
   }
 }