Sema: Record defaulted constraints in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -76,6 +76,8 @@ class SolverTrail {
     /// Recorded the mapping from a pack element expression to its parent
     /// pack expansion expression.
     RecordedPackEnvironment,
+    /// Record a defaulted constraint at a locator.
+    RecordedDefaultedConstraint,
   };
 
   /// A change made to the constraint system.
@@ -216,6 +218,9 @@ class SolverTrail {
     /// to its parent expansion expression.
     static Change recordedPackEnvironment(PackElementExpr *packElement);
 
+    /// Create a change that recorded a defaulted constraint at a locator.
+    static Change recordedDefaultedConstraint(ConstraintLocator *locator);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2412,10 +2412,25 @@ class ConstraintSystem {
 
 public:
   /// A map from argument expressions to their applied property wrapper expressions.
-  llvm::SmallMapVector<ASTNode, SmallVector<AppliedPropertyWrapper, 2>, 4> appliedPropertyWrappers;
+  llvm::SmallMapVector<ASTNode, SmallVector<AppliedPropertyWrapper, 2>, 4>
+      appliedPropertyWrappers;
 
   /// The locators of \c Defaultable constraints whose defaults were used.
-  llvm::SetVector<ConstraintLocator *> DefaultedConstraints;
+  llvm::DenseSet<ConstraintLocator *> DefaultedConstraints;
+
+  void recordDefaultedConstraint(ConstraintLocator *locator) {
+    bool inserted = DefaultedConstraints.insert(locator).second;
+    if (inserted) {
+      if (isRecordingChanges()) {
+        recordChange(SolverTrail::Change::recordedDefaultedConstraint(locator));
+      }
+    }
+  }
+
+  void removeDefaultedConstraint(ConstraintLocator *locator) {
+    bool erased = DefaultedConstraints.erase(locator);
+    ASSERT(erased);
+  }
 
   /// A cache that stores the @dynamicCallable required methods implemented by
   /// types.
@@ -2879,9 +2894,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c DefaultedConstraints.
-    unsigned numDefaultedConstraints;
-
     unsigned numAddedNodeTypes;
 
     unsigned numAddedKeyPathComponentTypes;

lib/Sema/CSBindings.cpp

@@ -2747,7 +2747,7 @@ bool TypeVariableBinding::attempt(ConstraintSystem &cs) const {
 
   // If this was from a defaultable binding note that.
   if (Binding.isDefaultableBinding()) {
-    cs.DefaultedConstraints.insert(srcLocator);
+    cs.recordDefaultedConstraint(srcLocator);
 
     // Fail if hole reporting fails.
     if (type->isPlaceholder() && reportHole())

lib/Sema/CSSolver.cpp

@@ -347,8 +347,8 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   }
 
   // Register the defaulted type variables.
-  DefaultedConstraints.insert(solution.DefaultedConstraints.begin(),
-                              solution.DefaultedConstraints.end());
+  for (auto *locator : solution.DefaultedConstraints)
+    recordDefaultedConstraint(locator);
 
   // Add the node types back.
   for (auto &nodeType : solution.nodeTypes) {
@@ -663,7 +663,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
 
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
-  numDefaultedConstraints = cs.DefaultedConstraints.size();
   numAddedNodeTypes = cs.addedNodeTypes.size();
   numAddedKeyPathComponentTypes = cs.addedKeyPathComponentTypes.size();
   numKeyPaths = cs.KeyPaths.size();
@@ -719,9 +718,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  // Remove any defaulted type variables.
-  truncate(cs.DefaultedConstraints, numDefaultedConstraints);
-
   // Remove any node types we registered.
   for (unsigned i :
            reverse(range(numAddedNodeTypes, cs.addedNodeTypes.size()))) {

lib/Sema/CSTrail.cpp

@@ -217,6 +217,14 @@ SolverTrail::Change::recordedPackEnvironment(PackElementExpr *packElement) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedDefaultedConstraint(ConstraintLocator *locator) {
+  Change result;
+  result.Kind = ChangeKind::RecordedDefaultedConstraint;
+  result.Locator = locator;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -301,6 +309,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedPackEnvironment:
     cs.removePackEnvironment(ElementExpr);
     break;
+
+  case ChangeKind::RecordedDefaultedConstraint:
+    cs.removeDefaultedConstraint(Locator);
+    break;
   }
 }
 
@@ -458,6 +470,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
   case ChangeKind::RecordedPackEnvironment:
     out << "(recorded pack environment)\n";
     break;
+
+  case ChangeKind::RecordedDefaultedConstraint:
+    out << "(recorded defaulted constraint at ";
+    Locator->dump(&cs.getASTContext().SourceMgr, out);
+    out << ")\n";
+    break;
   }
 }