Sema: Rework change recording in PotentialBindings::retract()

slavapestov · slavapestov · commit a36f63a86bed · 2025-02-04T17:53:28.000-05:00
Instead of making an undo() do an infer(), let's record fine-grained
changes about what was retracted, and directly re-insert the same
elements into the data structures.
diff --git a/include/swift/Sema/CSBindings.h b/include/swift/Sema/CSBindings.h
@@ -71,15 +71,12 @@ enum class LiteralBindingKind : uint8_t {
 /// along with information that can be used to construct related
 /// bindings, e.g., the supertypes of a given type.
 struct PotentialBinding {
-  friend class BindingSet;
-
   /// The type to which the type variable can be bound.
   Type BindingType;
 
   /// The kind of bindings permitted.
   AllowedBindingKind Kind;
 
-protected:
   /// The source of the type information.
   ///
   /// Determines whether this binding represents a "hole" in
@@ -91,7 +88,6 @@ struct PotentialBinding {
                    PointerUnion<Constraint *, ConstraintLocator *> source)
       : BindingType(type), Kind(kind), BindingSource(source) {}
 
-public:
   PotentialBinding(Type type, AllowedBindingKind kind, Constraint *source)
       : PotentialBinding(
             type, kind,
diff --git a/include/swift/Sema/CSTrail.def b/include/swift/Sema/CSTrail.def
@@ -38,6 +38,10 @@
 #define GRAPH_NODE_CHANGE(Name) CHANGE(Name)
 #endif
 
+#ifndef BINDING_RELATION_CHANGE
+#define BINDING_RELATION_CHANGE(Name) CHANGE(Name)
+#endif
+
 #ifndef SCORE_CHANGE
 #define SCORE_CHANGE(Name) CHANGE(Name)
 #endif
@@ -73,6 +77,12 @@ GRAPH_NODE_CHANGE(AddedConstraint)
 GRAPH_NODE_CHANGE(RemovedConstraint)
 GRAPH_NODE_CHANGE(InferredBindings)
 GRAPH_NODE_CHANGE(RetractedBindings)
+GRAPH_NODE_CHANGE(RetractedDelayedBy)
+
+BINDING_RELATION_CHANGE(RetractedAdjacentVar)
+BINDING_RELATION_CHANGE(RetractedSubtypeOf)
+BINDING_RELATION_CHANGE(RetractedSupertypeOf)
+BINDING_RELATION_CHANGE(RetractedEquivalentTo)
 
 SCORE_CHANGE(IncreasedScore)
 SCORE_CHANGE(DecreasedScore)
@@ -96,14 +106,16 @@ CHANGE(RecordedPotentialThrowSite)
 CHANGE(RecordedIsolatedParam)
 CHANGE(RecordedKeyPath)
 CHANGE(RetiredConstraint)
+CHANGE(RetractedBinding)
 
-LAST_CHANGE(RetiredConstraint)
+LAST_CHANGE(RetractedBinding)
 
 #undef LOCATOR_CHANGE
 #undef EXPR_CHANGE
 #undef CLOSURE_CHANGE
 #undef CONSTRAINT_CHANGE
 #undef GRAPH_NODE_CHANGE
+#undef BINDING_RELATION_CHANGE
 #undef SCORE_CHANGE
 #undef LAST_CHANGE
 #undef CHANGE
diff --git a/include/swift/Sema/CSTrail.h b/include/swift/Sema/CSTrail.h
@@ -38,6 +38,10 @@ namespace constraints {
 class Constraint;
 struct SyntacticElementTargetKey;
 
+namespace inference {
+struct PotentialBinding;
+}
+
 class SolverTrail {
 public:
 
@@ -89,6 +93,12 @@ class SolverTrail {
         TypeVariableType *OtherTypeVar;
       } Relation;
 
+      struct {
+        TypeVariableType *TypeVar;
+        TypeVariableType *OtherTypeVar;
+        Constraint *Constraint;
+      } BindingRelation;
+
       struct {
         /// The type variable being updated.
         TypeVariableType *TypeVar;
@@ -128,6 +138,15 @@ class SolverTrail {
         Constraint *Constraint;
       } Retiree;
 
+      struct {
+        TypeVariableType *TypeVar;
+
+        /// These two fields together with 'Options' above store the contents
+        /// of a PotentialBinding.
+        Type BindingType;
+        PointerUnion<Constraint *, ConstraintLocator *> BindingSource;
+      } Binding;
+
       ConstraintFix *TheFix;
       ConstraintLocator *TheLocator;
       PackExpansionType *TheExpansion;
@@ -155,6 +174,9 @@ class SolverTrail {
 #define SCORE_CHANGE(Name) static Change Name(ScoreKind kind, unsigned value);
 #define GRAPH_NODE_CHANGE(Name) static Change Name(TypeVariableType *typeVar, \
                                                    Constraint *constraint);
+#define BINDING_RELATION_CHANGE(Name) static Change Name(TypeVariableType *typeVar, \
+                                                         TypeVariableType *otherTypeVar, \
+                                                         Constraint *constraint);
 #include "swift/Sema/CSTrail.def"
 
     /// Create a change that added a type variable.
@@ -226,6 +248,11 @@ class SolverTrail {
     static Change RetiredConstraint(llvm::ilist<Constraint>::iterator where,
                                     Constraint *constraint);
 
+    /// Create a change that removed a binding from a type variable's potential
+    /// bindings.
+    static Change RetractedBinding(TypeVariableType *typeVar,
+                                   inference::PotentialBinding binding);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///
diff --git a/lib/Sema/CSBindings.cpp b/lib/Sema/CSBindings.cpp
@@ -2073,8 +2073,10 @@ void PotentialBindings::retract(ConstraintSystem &CS,
   if (!Constraints.remove(constraint))
     return;
 
+  bool recordingChanges = CS.solverState && !CS.solverState->Trail.isUndoActive();
+
   // Record the change, if there are active scopes.
-  if (CS.solverState && !CS.solverState->Trail.isUndoActive())
+  if (recordingChanges)
     CS.recordChange(SolverTrail::Change::RetractedBindings(TypeVar, constraint));
 
   LLVM_DEBUG(
@@ -2085,39 +2087,61 @@ void PotentialBindings::retract(ConstraintSystem &CS,
 
   Bindings.erase(
       llvm::remove_if(Bindings,
-                      [&constraint](const PotentialBinding &binding) {
-                        return binding.getSource() == constraint;
+                      [&](const PotentialBinding &binding) {
+                        if (binding.getSource() == constraint) {
+                          if (recordingChanges) {
+                            CS.recordChange(SolverTrail::Change::RetractedBinding(
+                                TypeVar, binding));
+                          }
+                          return true;
+                        }
+                        return false;
                       }),
       Bindings.end());
 
   DelayedBy.erase(
       llvm::remove_if(DelayedBy,
-                      [&constraint](Constraint *existing) {
-                        return existing == constraint;
+                      [&](Constraint *existing) {
+                        if (existing == constraint) {
+                          if (recordingChanges) {
+                            CS.recordChange(SolverTrail::Change::RetractedDelayedBy(
+                                TypeVar, constraint));
+                          }
+                          return true;
+                        }
+                        return false;
                       }),
       DelayedBy.end());
 
-  auto hasMatchingSource =
-      [&constraint](
-          const std::pair<TypeVariableType *, Constraint *> &adjacency) {
-        return adjacency.second == constraint;
-      };
+#define CALLBACK(ChangeKind)                                                   \
+  [&](std::pair<TypeVariableType *, Constraint *> pair) {                      \
+    if (pair.second == constraint) {                                           \
+      if (recordingChanges) {                                                  \
+        CS.recordChange(SolverTrail::Change::ChangeKind(                       \
+            TypeVar, pair.first, pair.second));                                \
+      }                                                                        \
+      return true;                                                             \
+    }                                                                          \
+    return false;                                                              \
+  }
 
   AdjacentVars.erase(
-    llvm::remove_if(AdjacentVars, hasMatchingSource),
+    llvm::remove_if(AdjacentVars, CALLBACK(RetractedAdjacentVar)),
     AdjacentVars.end());
 
   SubtypeOf.erase(
-    llvm::remove_if(SubtypeOf, hasMatchingSource),
+    llvm::remove_if(SubtypeOf, CALLBACK(RetractedSubtypeOf)),
     SubtypeOf.end());
 
   SupertypeOf.erase(
-    llvm::remove_if(SupertypeOf, hasMatchingSource),
+    llvm::remove_if(SupertypeOf, CALLBACK(RetractedSupertypeOf)),
     SupertypeOf.end());
 
   EquivalentTo.erase(
-    llvm::remove_if(EquivalentTo, hasMatchingSource),
+    llvm::remove_if(EquivalentTo, CALLBACK(RetractedEquivalentTo)),
     EquivalentTo.end());
+
+#undef CALLBACK
 }
 
 void PotentialBindings::reset() {
diff --git a/lib/Sema/CSTrail.cpp b/lib/Sema/CSTrail.cpp
@@ -90,6 +90,18 @@ SolverTrail::~SolverTrail() {
     result.TheConstraint.Constraint = constraint; \
     return result; \
   }
+#define BINDING_RELATION_CHANGE(Name) \
+  SolverTrail::Change \
+  SolverTrail::Change::Name(TypeVariableType *typeVar, \
+                            TypeVariableType *otherTypeVar, \
+                            Constraint *constraint) { \
+    Change result; \
+    result.Kind = ChangeKind::Name; \
+    result.BindingRelation.TypeVar = typeVar; \
+    result.BindingRelation.OtherTypeVar = otherTypeVar; \
+    result.BindingRelation.Constraint = constraint; \
+    return result; \
+  }
 #define SCORE_CHANGE(Name) \
   SolverTrail::Change \
   SolverTrail::Change::Name(ScoreKind kind, unsigned value) { \
@@ -306,6 +318,19 @@ SolverTrail::Change::RetiredConstraint(ConstraintList::iterator where,
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::RetractedBinding(TypeVariableType *typeVar,
+                                      inference::PotentialBinding binding) {
+  Change result;
+  result.Kind = ChangeKind::RetractedBinding;
+  result.Binding.TypeVar = typeVar;
+  result.Binding.BindingType = binding.BindingType;
+  result.Binding.BindingSource = binding.BindingSource;
+  result.Options = unsigned(binding.Kind);
+
+  return result;
+}
+
 SyntacticElementTargetKey
 SolverTrail::Change::getSyntacticElementTargetKey() const {
   ASSERT(Kind == ChangeKind::RecordedTarget);
@@ -359,11 +384,9 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     cg.addConstraint(TheConstraint.TypeVar, TheConstraint.Constraint);
     break;
 
-  case ChangeKind::ExtendedEquivalenceClass: {
-    auto &node = cg[EquivClass.TypeVar];
-    node.truncateEquivalenceClass(EquivClass.PrevSize);
+  case ChangeKind::ExtendedEquivalenceClass:
+    cg[EquivClass.TypeVar].truncateEquivalenceClass(EquivClass.PrevSize);
     break;
-   }
 
   case ChangeKind::RelatedTypeVariables:
     cg.unrelateTypeVariables(Relation.TypeVar, Relation.OtherTypeVar);
@@ -373,9 +396,12 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     cg.retractBindings(TheConstraint.TypeVar, TheConstraint.Constraint);
     break;
 
-  case ChangeKind::RetractedBindings:
-    cg.inferBindings(TheConstraint.TypeVar, TheConstraint.Constraint);
+  case ChangeKind::RetractedBindings: {
+    auto &bindings = cg[TheConstraint.TypeVar].getPotentialBindings();
+    bool inserted = bindings.Constraints.insert(TheConstraint.Constraint);
+    ASSERT(inserted);
     break;
+  }
 
   case ChangeKind::UpdatedTypeVariable:
     Update.TypeVar->getImpl().setRawOptions(Options);
@@ -494,6 +520,45 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     cs.InactiveConstraints.insert(Retiree.Where,
                                   Retiree.Constraint);
     break;
+
+  case ChangeKind::RetractedDelayedBy:
+    cg[TheConstraint.TypeVar].getPotentialBindings()
+        .DelayedBy.push_back(TheConstraint.Constraint);
+    break;
+
+  case ChangeKind::RetractedAdjacentVar:
+    cg[BindingRelation.TypeVar].getPotentialBindings()
+        .AdjacentVars.emplace_back(BindingRelation.OtherTypeVar,
+                                   BindingRelation.Constraint);
+    break;
+
+  case ChangeKind::RetractedSubtypeOf:
+    cg[BindingRelation.TypeVar].getPotentialBindings()
+        .SubtypeOf.emplace_back(BindingRelation.OtherTypeVar,
+                                BindingRelation.Constraint);
+    break;
+
+  case ChangeKind::RetractedSupertypeOf:
+    cg[BindingRelation.TypeVar].getPotentialBindings()
+        .SupertypeOf.emplace_back(BindingRelation.OtherTypeVar,
+                                  BindingRelation.Constraint);
+    break;
+
+  case ChangeKind::RetractedEquivalentTo:
+    cg[BindingRelation.TypeVar].getPotentialBindings()
+        .EquivalentTo.emplace_back(BindingRelation.OtherTypeVar,
+                                   BindingRelation.Constraint);
+    break;
+
+  case ChangeKind::RetractedBinding: {
+    PotentialBinding binding(Binding.BindingType,
+                             AllowedBindingKind(Options),
+                             Binding.BindingSource);
+
+    auto &bindings = cg[BindingRelation.TypeVar].getPotentialBindings();
+    bindings.Bindings.push_back(binding);
+    break;
+  }
   }
 }
 
@@ -541,6 +606,17 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
       TheConstraint.TypeVar->print(out, PO); \
       out << ")\n"; \
       break;
+#define BINDING_RELATION_CHANGE(Name) \
+    case ChangeKind::Name: \
+      out << "(" << #Name << " "; \
+      BindingRelation.Constraint->print(out, &cs.getASTContext().SourceMgr, \
+                                        indent + 2); \
+      out << " on type variable "; \
+      BindingRelation.TypeVar->print(out, PO); \
+      out << " and "; \
+      BindingRelation.OtherTypeVar->print(out, PO); \
+      out << ")\n"; \
+      break;
 #define SCORE_CHANGE(Name) \
     case ChangeKind::Name: \
       out << "(" << #Name << " "; \
@@ -692,6 +768,14 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
                               indent + 2);
     out << ")\n";
     break;
+
+  case ChangeKind::RetractedBinding:
+    out << "(RetractedBinding ";
+    Binding.TypeVar->print(out, PO);
+    out << " with type ";
+    Binding.BindingType->print(out, PO);
+    out << " and kind " << Options << ")\n";
+    break;
   }
 }
 
