Merge pull request #76742 from slavapestov/cstrail-part-1

slavapestov · web-flow · commit 106ab41d1d2b · 2024-10-01T12:36:13.000-04:00
Sema: Continue unwinding SolverScope
diff --git a/include/swift/Sema/CSBindings.h b/include/swift/Sema/CSBindings.h
@@ -223,6 +223,9 @@ struct PotentialBindings {
 
   TypeVariableType *TypeVar;
 
+  /// The set of all constraints that have been added via infer().
+  llvm::SmallPtrSet<Constraint *, 2> Constraints;
+
   /// The set of potential bindings.
   llvm::SmallVector<PotentialBinding, 4> Bindings;
 
diff --git a/include/swift/Sema/CSTrail.h b/include/swift/Sema/CSTrail.h
@@ -37,18 +37,20 @@ class SolverTrail {
 
   /// The kind of change made to the graph.
   enum class ChangeKind {
-    /// Added a type variable to the constraint graph.
+    /// Added a new vertex to the constraint graph.
     AddedTypeVariable,
     /// Added a new constraint to the constraint graph.
     AddedConstraint,
     /// Removed an existing constraint from the constraint graph.
     RemovedConstraint,
     /// Extended the equivalence class of a type variable in the constraint graph.
     ExtendedEquivalenceClass,
-    /// Added a fixed binding for a type variable in the constraint graph.
-    BoundTypeVariable,
-    /// Introduced a type variable's fixed type to inference.
-    IntroducedToInference,
+    /// Added a new edge in the constraint graph.
+    RelatedTypeVariables,
+    /// Inferred potential bindings from a constraint.
+    InferredBindings,
+    /// Retracted potential bindings from a constraint.
+    RetractedBindings,
     /// Set the fixed type or parent and flags for a type variable.
     UpdatedTypeVariable,
   };
@@ -64,7 +66,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.
@@ -75,12 +84,12 @@ class SolverTrail {
       } EquivClass;
 
       struct {
-        /// The type variable being bound to a fixed type.
+        /// The first type variable.
         TypeVariableType *TypeVar;
 
-        /// The fixed type to which the type variable was bound.
-        TypeBase *FixedType;
-      } Binding;
+        /// The second type variable.
+        TypeVariableType *OtherTypeVar;
+      } Relation;
 
       struct {
         /// The type variable being updated.
@@ -100,20 +109,27 @@ 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,
                                            unsigned prevSize);
 
-    /// Create a change that bound a type variable to a fixed type.
-    static Change boundTypeVariable(TypeVariableType *typeVar, Type fixed);
+    /// Create a change that updated the references/referenced by sets of
+    /// a type variable pair.
+    static Change relatedTypeVariables(TypeVariableType *typeVar,
+                                       TypeVariableType *otherTypeVar);
+
+    /// Create a change that inferred bindings from a constraint.
+    static Change inferredBindings(TypeVariableType *typeVar,
+                                   Constraint *constraint);
 
-    /// Create a change that introduced a type variable to inference.
-    static Change introducedToInference(TypeVariableType *typeVar, Type fixed);
+    /// Create a change that retracted bindings from a constraint.
+    static Change retractedBindings(TypeVariableType *typeVar,
+                                    Constraint *constraint);
 
     /// Create a change that updated a type variable.
     static Change updatedTypeVariable(
diff --git a/include/swift/Sema/ConstraintGraph.h b/include/swift/Sema/ConstraintGraph.h
@@ -153,13 +153,6 @@ class ConstraintGraphNode {
   /// type variable has been bound to a valid type and solver can make progress.
   void introduceToInference(Type fixedType);
 
-  /// Opposite of \c introduceToInference(Type)
-  void retractFromInference(Type fixedType);
-
-  /// Drop all previously collected bindings and re-infer based on the
-  /// current set constraints associated with this equivalence class.
-  void resetBindingSet();
-
   /// Notify all of the type variables that have this one (or any member of
   /// its equivalence class) referenced in their fixed type.
   ///
@@ -259,9 +252,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 +415,29 @@ 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);
 
-  /// Unbind the given type variable from the given fixed type.
+  /// Remove an edge from the constraint graph.
+  ///
+  ///
+  /// Note that this it only meant to be called by SolverTrail::Change::undo().
+  void unrelateTypeVariables(TypeVariableType *typeVar,
+                             TypeVariableType *otherTypeVar);
+
+  /// Infer bindings from the given constraint.
   ///
-  /// Note that this change is not recorded and cannot be undone. Use with
-  /// caution.
-  void unbindTypeVariable(TypeVariableType *typeVar, Type fixedType);
+  /// Note that this it only meant to be called by SolverTrail::Change::undo().
+  void inferBindings(TypeVariableType *typeVar, Constraint *constraint);
 
-  /// Retract the given type variable from inference.
+  /// Retract bindings from the given constraint.
   ///
-  /// Note that this change is not recorded and cannot be undone. Use with
-  /// caution.
-  void retractFromInference(TypeVariableType *typeVar, Type fixedType);
+  /// Note that this it only meant to be called by SolverTrail::Change::undo().
+  void retractBindings(TypeVariableType *typeVar, Constraint *constraint);
 
   /// Perform edge contraction on the constraint graph, merging equivalence
   /// classes until a fixed point is reached.
diff --git a/lib/Sema/CSBindings.cpp b/lib/Sema/CSBindings.cpp
@@ -1769,6 +1769,13 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
 /// representative type variable, along with flags indicating whether
 /// those types should be opened.
 void PotentialBindings::infer(Constraint *constraint) {
+  if (!Constraints.insert(constraint).second)
+    return;
+
+  // Record the change, if there are active scopes.
+  if (CS.isRecordingChanges())
+    CS.recordChange(SolverTrail::Change::inferredBindings(TypeVar, constraint));
+
   switch (constraint->getKind()) {
   case ConstraintKind::Bind:
   case ConstraintKind::Equal:
@@ -1937,6 +1944,13 @@ void PotentialBindings::infer(Constraint *constraint) {
 }
 
 void PotentialBindings::retract(Constraint *constraint) {
+  if (!Constraints.erase(constraint))
+    return;
+
+  // Record the change, if there are active scopes.
+  if (CS.isRecordingChanges())
+    CS.recordChange(SolverTrail::Change::retractedBindings(TypeVar, constraint));
+
   Bindings.erase(
       llvm::remove_if(Bindings,
                       [&constraint](const PotentialBinding &binding) {
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;
 }
 
@@ -76,22 +80,32 @@ SolverTrail::Change::extendedEquivalenceClass(TypeVariableType *typeVar,
 }
 
 SolverTrail::Change
-SolverTrail::Change::boundTypeVariable(TypeVariableType *typeVar,
-                                       Type fixed) {
+SolverTrail::Change::relatedTypeVariables(TypeVariableType *typeVar,
+                                          TypeVariableType *otherTypeVar) {
   Change result;
-  result.Kind = ChangeKind::BoundTypeVariable;
-  result.Binding.TypeVar = typeVar;
-  result.Binding.FixedType = fixed.getPointer();
+  result.Kind = ChangeKind::RelatedTypeVariables;
+  result.Relation.TypeVar = typeVar;
+  result.Relation.OtherTypeVar = otherTypeVar;
   return result;
 }
 
 SolverTrail::Change
-SolverTrail::Change::introducedToInference(TypeVariableType *typeVar,
-                                           Type fixed) {
+SolverTrail::Change::inferredBindings(TypeVariableType *typeVar,
+                                     Constraint *constraint) {
   Change result;
-  result.Kind = ChangeKind::IntroducedToInference;
-  result.Binding.TypeVar = typeVar;
-  result.Binding.FixedType = fixed.getPointer();
+  result.Kind = ChangeKind::InferredBindings;
+  result.TheConstraint.TypeVar = typeVar;
+  result.TheConstraint.Constraint = constraint;
+  return result;
+}
+
+SolverTrail::Change
+SolverTrail::Change::retractedBindings(TypeVariableType *typeVar,
+                                       Constraint *constraint) {
+  Change result;
+  result.Kind = ChangeKind::RetractedBindings;
+  result.TheConstraint.TypeVar = typeVar;
+  result.TheConstraint.Constraint = constraint;
   return result;
 }
 
@@ -117,11 +131,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: {
@@ -130,12 +144,16 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     break;
    }
 
-  case ChangeKind::BoundTypeVariable:
-    cg.unbindTypeVariable(Binding.TypeVar, Binding.FixedType);
+  case ChangeKind::RelatedTypeVariables:
+    cg.unrelateTypeVariables(Relation.TypeVar, Relation.OtherTypeVar);
     break;
 
-  case ChangeKind::IntroducedToInference:
-    cg.retractFromInference(Binding.TypeVar, Binding.FixedType);
+  case ChangeKind::InferredBindings:
+    cg.retractBindings(TheConstraint.TypeVar, TheConstraint.Constraint);
+    break;
+
+  case ChangeKind::RetractedBindings:
+    cg.inferBindings(TheConstraint.TypeVar, TheConstraint.Constraint);
     break;
 
   case ChangeKind::UpdatedTypeVariable:
@@ -162,13 +180,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;
 
@@ -179,20 +203,30 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     break;
    }
 
-  case ChangeKind::BoundTypeVariable:
-    out << "(bound type variable ";
-    Binding.TypeVar->print(out, PO);
-    out << " to fixed type ";
-    Binding.FixedType->print(out, PO);
+  case ChangeKind::RelatedTypeVariables:
+    out << "(related type variable ";
+    Relation.TypeVar->print(out, PO);
+    out << " with ";
+    Relation.OtherTypeVar->print(out, PO);
     out << ")\n";
     break;
 
-  case ChangeKind::IntroducedToInference:
-    out << "(introduced type variable ";
-    Binding.TypeVar->print(out, PO);
-    out << " with fixed type ";
-    Binding.FixedType->print(out, PO);
-    out << " to inference)\n";
+  case ChangeKind::InferredBindings:
+    out << "(inferred bindings from ";
+    TheConstraint.Constraint->print(out, &cs.getASTContext().SourceMgr,
+                         indent + 2);
+    out << " for type variable ";
+    TheConstraint.TypeVar->print(out, PO);
+    out << ")\n";
+    break;
+
+  case ChangeKind::RetractedBindings:
+    out << "(retracted bindings from ";
+    TheConstraint.Constraint->print(out, &cs.getASTContext().SourceMgr,
+                                    indent + 2);
+    out << " for type variable ";
+    TheConstraint.TypeVar->print(out, PO);
+    out << ")\n";
     break;
 
   case ChangeKind::UpdatedTypeVariable:
diff --git a/lib/Sema/ConstraintGraph.cpp b/lib/Sema/ConstraintGraph.cpp
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
