Sema: Split up SolverTrail::Change::AddedConstraint and ::RemovedConstraint

slavapestov · slavapestov · commit b885c29f7159 · 2024-09-30T21:29:59.000-04:00
diff --git a/include/swift/Sema/CSTrail.h b/include/swift/Sema/CSTrail.h
@@ -64,7 +64,14 @@ class SolverTrail {
 
     union {
       TypeVariableType *TypeVar;
-      Constraint *TheConstraint;
+
+      struct {
+        /// The type variable we're adding or removing a constraint from.
+        TypeVariableType *TypeVar;
+
+        /// The constraint.
+        Constraint *Constraint;
+      } TheConstraint;
 
       struct {
         /// The type variable whose equivalence class was extended.
@@ -108,10 +115,10 @@ class SolverTrail {
     static Change addedTypeVariable(TypeVariableType *typeVar);
 
     /// Create a change that added a constraint.
-    static Change addedConstraint(Constraint *constraint);
+    static Change addedConstraint(TypeVariableType *typeVar, Constraint *constraint);
 
     /// Create a change that removed a constraint.
-    static Change removedConstraint(Constraint *constraint);
+    static Change removedConstraint(TypeVariableType *typeVar, Constraint *constraint);
 
     /// Create a change that extended an equivalence class.
     static Change extendedEquivalenceClass(TypeVariableType *typeVar,
diff --git a/include/swift/Sema/ConstraintGraph.h b/include/swift/Sema/ConstraintGraph.h
@@ -259,9 +259,15 @@ class ConstraintGraph {
   /// Add a new constraint to the graph.
   void addConstraint(Constraint *constraint);
 
+  /// Primitive form for SolverTrail::Change::undo().
+  void addConstraint(TypeVariableType *typeVar, Constraint *constraint);
+
   /// Remove a constraint from the graph.
   void removeConstraint(Constraint *constraint);
 
+  /// Primitive form for SolverTrail::Change::undo().
+  void removeConstraint(TypeVariableType *typeVar, Constraint *constraint);
+
   /// Merge the two nodes for the two given type variables.
   ///
   /// The type variables must actually have been merged already; this
@@ -416,25 +422,23 @@ class ConstraintGraph {
 private:
   /// Remove the node corresponding to the given type variable.
   ///
-  /// This operation assumes that the any constraints that refer to
-  /// this type variable have been or will be removed before other
-  /// graph queries are performed.
+  /// This operation assumes that the any constraints that refer to this type
+  /// variable have been or will be removed before other graph queries are
+  /// performed.
   ///
-  /// Note that this change is not recorded and cannot be undone. Use with
-  /// caution.
+  /// Note that this it only meant to be called by SolverTrail::Change::undo().
   void removeNode(TypeVariableType *typeVar);
 
   /// Remove an edge from the constraint graph.
   ///
-  /// Note that this change is not recorded and cannot be undone. Use with
-  /// caution.
+  ///
+  /// Note that this it only meant to be called by SolverTrail::Change::undo().
   void unrelateTypeVariables(TypeVariableType *typeVar,
                              TypeVariableType *otherTypeVar);
 
   /// Retract the given type variable from inference.
   ///
-  /// Note that this change is not recorded and cannot be undone. Use with
-  /// caution.
+  /// Note that this it only meant to be called by SolverTrail::Change::undo().
   void retractFromInference(TypeVariableType *typeVar, Type fixedType);
 
   /// Perform edge contraction on the constraint graph, merging equivalence
diff --git a/lib/Sema/CSTrail.cpp b/lib/Sema/CSTrail.cpp
@@ -50,18 +50,22 @@ SolverTrail::Change::addedTypeVariable(TypeVariableType *typeVar) {
 }
 
 SolverTrail::Change
-SolverTrail::Change::addedConstraint(Constraint *constraint) {
+SolverTrail::Change::addedConstraint(TypeVariableType *typeVar,
+                                     Constraint *constraint) {
   Change result;
   result.Kind = ChangeKind::AddedConstraint;
-  result.TheConstraint = constraint;
+  result.TheConstraint.TypeVar = typeVar;
+  result.TheConstraint.Constraint = constraint;
   return result;
 }
 
 SolverTrail::Change
-SolverTrail::Change::removedConstraint(Constraint *constraint) {
+SolverTrail::Change::removedConstraint(TypeVariableType *typeVar,
+                                       Constraint *constraint) {
   Change result;
   result.Kind = ChangeKind::RemovedConstraint;
-  result.TheConstraint = constraint;
+  result.TheConstraint.TypeVar = typeVar;
+  result.TheConstraint.Constraint = constraint;
   return result;
 }
 
@@ -117,11 +121,11 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     break;
 
   case ChangeKind::AddedConstraint:
-    cg.removeConstraint(TheConstraint);
+    cg.removeConstraint(TheConstraint.TypeVar, TheConstraint.Constraint);
     break;
 
   case ChangeKind::RemovedConstraint:
-    cg.addConstraint(TheConstraint);
+    cg.addConstraint(TheConstraint.TypeVar, TheConstraint.Constraint);
     break;
 
   case ChangeKind::ExtendedEquivalenceClass: {
@@ -162,13 +166,19 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
 
   case ChangeKind::AddedConstraint:
     out << "(added constraint ";
-    TheConstraint->print(out, &cs.getASTContext().SourceMgr, indent + 2);
+    TheConstraint.Constraint->print(out, &cs.getASTContext().SourceMgr,
+                         indent + 2);
+    out << " to type variable ";
+    TheConstraint.TypeVar->print(out, PO);
     out << ")\n";
     break;
 
   case ChangeKind::RemovedConstraint:
     out << "(removed constraint ";
-    TheConstraint->print(out, &cs.getASTContext().SourceMgr, indent + 2);
+    TheConstraint.Constraint->print(out, &cs.getASTContext().SourceMgr,
+                                    indent + 2);
+    out << " from type variable ";
+    TheConstraint.TypeVar->print(out, PO);
     out << ")\n";
     break;
 
diff --git a/lib/Sema/ConstraintGraph.cpp b/lib/Sema/ConstraintGraph.cpp
@@ -155,16 +155,6 @@ void ConstraintGraphNode::addConstraint(Constraint *constraint) {
   assert(ConstraintIndex.count(constraint) == 0 && "Constraint re-insertion");
   ConstraintIndex[constraint] = Constraints.size();
   Constraints.push_back(constraint);
-
-  {
-    introduceToInference(constraint);
-
-    if (isUsefulForReferencedVars(constraint)) {
-      notifyReferencedVars([&](ConstraintGraphNode &referencedVar) {
-        referencedVar.introduceToInference(constraint);
-      });
-    }
-  }
 }
 
 void ConstraintGraphNode::removeConstraint(Constraint *constraint) {
@@ -176,16 +166,6 @@ void ConstraintGraphNode::removeConstraint(Constraint *constraint) {
   ConstraintIndex.erase(pos);
   assert(Constraints[index] == constraint && "Mismatched constraint");
 
-  {
-    retractFromInference(constraint);
-
-    if (isUsefulForReferencedVars(constraint)) {
-      notifyReferencedVars([&](ConstraintGraphNode &referencedVar) {
-        referencedVar.retractFromInference(constraint);
-      });
-    }
-  }
-
   // If this is the last constraint, just pop it off the list and we're done.
   unsigned lastIndex = Constraints.size()-1;
   if (index == lastIndex) {
@@ -471,48 +451,97 @@ void ConstraintGraph::addConstraint(Constraint *constraint) {
   // For the nodes corresponding to each type variable...
   auto referencedTypeVars = constraint->getTypeVariables();
   for (auto typeVar : referencedTypeVars) {
+    // Record the change, if there are active scopes.
+    if (CS.isRecordingChanges())
+      CS.recordChange(SolverTrail::Change::addedConstraint(typeVar, constraint));
+
+    addConstraint(typeVar, constraint);
+  }
+
+  // If the constraint doesn't reference any type variables, it's orphaned;
+  // track it as such.
+  if (referencedTypeVars.empty()) {
+    // Record the change, if there are active scopes.
+    if (CS.isRecordingChanges())
+      CS.recordChange(SolverTrail::Change::addedConstraint(nullptr, constraint));
+
+    addConstraint(nullptr, constraint);
+  }
+}
+
+void ConstraintGraph::addConstraint(TypeVariableType *typeVar,
+                                    Constraint *constraint) {
+  if (typeVar) {
     // Find the node for this type variable.
     auto &node = (*this)[typeVar];
 
     // Note the constraint within the node for that type variable.
     node.addConstraint(constraint);
+
+    node.introduceToInference(constraint);
+
+    if (isUsefulForReferencedVars(constraint)) {
+      node.notifyReferencedVars([&](ConstraintGraphNode &referencedVar) {
+        referencedVar.introduceToInference(constraint);
+      });
+    }
+
+    return;
   }
 
   // If the constraint doesn't reference any type variables, it's orphaned;
   // track it as such.
-  if (referencedTypeVars.empty()) {
-    OrphanedConstraints.push_back(constraint);
-  }
-
-  // Record the change, if there are active scopes.
-  if (CS.isRecordingChanges())
-    CS.recordChange(SolverTrail::Change::addedConstraint(constraint));
+  OrphanedConstraints.push_back(constraint);
 }
 
 void ConstraintGraph::removeConstraint(Constraint *constraint) {
   // For the nodes corresponding to each type variable...
   auto referencedTypeVars = constraint->getTypeVariables();
   for (auto typeVar : referencedTypeVars) {
+    // Record the change, if there are active scopes.
+    if (CS.isRecordingChanges())
+      CS.recordChange(SolverTrail::Change::removedConstraint(typeVar, constraint));
+
+    removeConstraint(typeVar, constraint);
+  }
+
+  // If this is an orphaned constraint, remove it from the list.
+  if (referencedTypeVars.empty()) {
+    // Record the change, if there are active scopes.
+    if (CS.isRecordingChanges())
+      CS.recordChange(SolverTrail::Change::removedConstraint(nullptr, constraint));
+
+    removeConstraint(nullptr, constraint);
+  }
+}
+
+void ConstraintGraph::removeConstraint(TypeVariableType *typeVar,
+                                       Constraint *constraint) {
+  if (typeVar) {
     // Find the node for this type variable.
     auto &node = (*this)[typeVar];
 
+    node.retractFromInference(constraint);
+
+    if (isUsefulForReferencedVars(constraint)) {
+      node.notifyReferencedVars([&](ConstraintGraphNode &referencedVar) {
+        referencedVar.retractFromInference(constraint);
+      });
+    }
+
     // Remove the constraint.
     node.removeConstraint(constraint);
-  }
 
-  // If this is an orphaned constraint, remove it from the list.
-  if (referencedTypeVars.empty()) {
-    auto known = std::find(OrphanedConstraints.begin(),
-                           OrphanedConstraints.end(),
-                           constraint);
-    assert(known != OrphanedConstraints.end() && "missing orphaned constraint");
-    *known = OrphanedConstraints.back();
-    OrphanedConstraints.pop_back();
+    return;
   }
 
-  // Record the change, if there are active scopes.
-  if (CS.isRecordingChanges())
-    CS.recordChange(SolverTrail::Change::removedConstraint(constraint));
+  // If this is an orphaned constraint, remove it from the list.
+  auto known = std::find(OrphanedConstraints.begin(),
+                         OrphanedConstraints.end(),
+                         constraint);
+  assert(known != OrphanedConstraints.end() && "missing orphaned constraint");
+  *known = OrphanedConstraints.back();
+  OrphanedConstraints.pop_back();
 }
 
 void ConstraintGraph::mergeNodes(TypeVariableType *typeVar1, 
