Sema: Factor out SolverTrail from ConstraintGraph

slavapestov · slavapestov · commit 4cc27c9b206c · 2024-09-26T23:15:26.000-04:00
diff --git a/include/swift/Sema/CSTrail.h b/include/swift/Sema/CSTrail.h
@@ -0,0 +1,138 @@
+//===--- CSTrail.h - Constraint Solver Trail --------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the \c SolverTrail class, which records the decisions taken
+// while attempting to find a solution.
+//
+//===----------------------------------------------------------------------===//
+#ifndef SWIFT_SEMA_CSTRAIL_H
+#define SWIFT_SEMA_CSTRAIL_H
+
+#include <vector>
+
+namespace llvm {
+class raw_ostream;
+}
+
+namespace swift {
+
+class TypeBase;
+class TypeVariableType;
+
+namespace constraints {
+
+class Constraint;
+class ConstraintGraphScope;
+
+class SolverTrail {
+public:
+
+  /// The kind of change made to the graph.
+  enum class ChangeKind {
+    /// Added a type variable 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,
+  };
+
+  /// A change made to the constraint system.
+  ///
+  /// Each change can be undone (once, and in reverse order) by calling the
+  /// undo() method.
+  class Change {
+  public:
+    /// The kind of change.
+    ChangeKind Kind;
+
+    union {
+      TypeVariableType *TypeVar;
+      Constraint *TheConstraint;
+
+      struct {
+        /// The type variable whose equivalence class was extended.
+        TypeVariableType *TypeVar;
+
+        /// The previous size of the equivalence class.
+        unsigned PrevSize;
+      } EquivClass;
+
+      struct {
+        /// The type variable being bound to a fixed type.
+        TypeVariableType *TypeVar;
+
+        /// The fixed type to which the type variable was bound.
+        TypeBase *FixedType;
+      } Binding;
+    };
+
+    Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
+
+    /// Create a change that added a type variable.
+    static Change addedTypeVariable(TypeVariableType *typeVar);
+
+    /// Create a change that added a constraint.
+    static Change addedConstraint(Constraint *constraint);
+
+    /// Create a change that removed a constraint.
+    static Change removedConstraint(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);
+
+    /// Undo this change, reverting the constraint graph to the state it
+    /// had prior to this change.
+    ///
+    /// Changes must be undone in stack order.
+    void undo(ConstraintSystem &cs);
+  };
+
+  SolverTrail(ConstraintSystem &cs) : CS(cs) {}
+
+  ~SolverTrail();
+
+  SolverTrail(const SolverTrail &) = delete;
+  SolverTrail &operator=(const SolverTrail &) = delete;
+
+  void recordChange(Change change);
+
+  void dumpActiveScopeChanges(llvm::raw_ostream &out,
+                              unsigned indent = 0);
+
+  unsigned size() const {
+    return Changes.size();
+  }
+
+  void undo(unsigned toIndex);
+
+private:
+  ConstraintSystem &CS;
+
+  /// The list of changes made to this constraint system.
+  std::vector<Change> Changes;
+
+  friend class ConstraintGraphScope;
+};
+
+} // namespace constraints
+} // namespace swift
+
+#endif // SWIFT_SEMA_CSTRAIL_H
diff --git a/include/swift/Sema/ConstraintGraph.h b/include/swift/Sema/ConstraintGraph.h
@@ -216,6 +216,7 @@ class ConstraintGraphNode {
   friend class ConstraintGraph;
   friend class ConstraintSystem;
   friend class TypeVariableBinding;
+  friend class SolverTrail;
 };
 
 /// A graph that describes the relationships among the various type variables
@@ -387,7 +388,6 @@ class ConstraintGraph {
   /// Print the graph.
   void print(ArrayRef<TypeVariableType *> typeVars, llvm::raw_ostream &out);
   void dump(llvm::raw_ostream &out);
-  void dumpActiveScopeChanges(llvm::raw_ostream &out, unsigned indent = 0);
 
   // FIXME: Potentially side-effectful.
   SWIFT_DEBUG_HELPER(void dump());
@@ -440,86 +440,12 @@ class ConstraintGraph {
   /// to contract constraint graph edges.
   void incrementConstraintsPerContractionCounter();
 
-  /// The kind of change made to the graph.
-  enum class ChangeKind {
-    /// Added a type variable.
-    AddedTypeVariable,
-    /// Added a new constraint.
-    AddedConstraint,
-    /// Removed an existing constraint
-    RemovedConstraint,
-    /// Extended the equivalence class of a type variable.
-    ExtendedEquivalenceClass,
-    /// Added a fixed binding for a type variable.
-    BoundTypeVariable,
-  };
-
-  /// A change made to the constraint graph.
-  ///
-  /// Each change can be undone (once, and in reverse order) by calling the
-  /// undo() method.
-  class Change {
-  public:
-    /// The kind of change.
-    ChangeKind Kind;
-
-    union {
-      TypeVariableType *TypeVar;
-      Constraint *TheConstraint;
-
-      struct {
-        /// The type variable whose equivalence class was extended.
-        TypeVariableType *TypeVar;
-
-        /// The previous size of the equivalence class.
-        unsigned PrevSize;
-      } EquivClass;
-
-      struct {
-        /// The type variable being bound to a fixed type.
-        TypeVariableType *TypeVar;
-
-        /// The fixed type to which the type variable was bound.
-        TypeBase *FixedType;
-      } Binding;
-    };
-
-    Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
-
-    /// Create a change that added a type variable.
-    static Change addedTypeVariable(TypeVariableType *typeVar);
-
-    /// Create a change that added a constraint.
-    static Change addedConstraint(Constraint *constraint);
-
-    /// Create a change that removed a constraint.
-    static Change removedConstraint(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);
-
-    /// Undo this change, reverting the constraint graph to the state it
-    /// had prior to this change.
-    ///
-    /// Changes must be undone in stack order.
-    void undo(ConstraintGraph &cg);
-  };
-
   /// The currently active scope, or null if we aren't tracking changes made
   /// to the constraint graph.
   ConstraintGraphScope *ActiveScope = nullptr;
 
-  /// The set of changes made to this constraint graph.
-  ///
-  /// As the constraint graph is extended and mutated, additional changes are
-  /// introduced into this vector. Each scope
-  llvm::SmallVector<Change, 4> Changes;
-
   friend class ConstraintGraphScope;
+  friend class SolverTrail;
 };
 
 } // end namespace constraints
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -32,6 +32,7 @@
 #include "swift/Basic/OptionSet.h"
 #include "swift/Sema/CSBindings.h"
 #include "swift/Sema/CSFix.h"
+#include "swift/Sema/CSTrail.h"
 #include "swift/Sema/Constraint.h"
 #include "swift/Sema/ConstraintGraph.h"
 #include "swift/Sema/ConstraintGraphScope.h"
@@ -63,8 +64,6 @@ enum class TypeCheckExprFlags;
 
 namespace constraints {
 
-class ConstraintGraph;
-class ConstraintGraphNode;
 class ConstraintSystem;
 class SyntacticElementTarget;
 
@@ -2533,6 +2532,10 @@ class ConstraintSystem {
     /// processing this constraint system.
     SavedTypeVariableBindings savedBindings;
 
+    /// A log of changes to the constraint system, representing the
+    /// current path being explored in the solution space.
+    SolverTrail Trail;
+
      /// The best solution computed so far.
     std::optional<Score> BestScore;
 
@@ -2681,6 +2684,10 @@ class ConstraintSystem {
       favoredConstraints.push_back(constraint);
     }
 
+    void recordChange(SolverTrail::Change change) {
+      Trail.recordChange(change);
+    }
+
   private:
     /// The list of constraints that have been retired along the
     /// current path, this list is used in LIFO fashion when
diff --git a/lib/Sema/CMakeLists.txt b/lib/Sema/CMakeLists.txt
@@ -11,6 +11,7 @@ add_swift_host_library(swiftSema STATIC
   CSSimplify.cpp
   CSSolver.cpp
   CSStep.cpp
+  CSTrail.cpp
   CSFix.cpp
   CSDiagnostics.cpp
   CodeSynthesis.cpp
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -568,7 +568,7 @@ void truncate(llvm::SetVector<V> &vector, unsigned newSize) {
 
 ConstraintSystem::SolverState::SolverState(
     ConstraintSystem &cs, FreeTypeVariableBinding allowFreeTypeVariables)
-    : CS(cs), AllowFreeTypeVariables(allowFreeTypeVariables) {
+    : CS(cs), AllowFreeTypeVariables(allowFreeTypeVariables), Trail(cs) {
   assert(!CS.solverState &&
          "Constraint system should not already have solver state!");
   CS.solverState = this;
diff --git a/lib/Sema/CSStep.h b/lib/Sema/CSStep.h
@@ -550,9 +550,8 @@ template <typename P> class BindingStep : public SolverStep {
           ActiveChoice.emplace(std::move(scope), *choice);
 
           if (CS.isDebugMode()) {
-            auto &log = llvm::errs();
-            auto &CG = CS.getConstraintGraph();
-            CG.dumpActiveScopeChanges(log, CS.solverState->getCurrentIndent());
+            CS.solverState->Trail.dumpActiveScopeChanges(
+              llvm::errs(), CS.solverState->getCurrentIndent());
           }
           
           return suspend(std::make_unique<SplitterStep>(CS, Solutions));
diff --git a/lib/Sema/CSTrail.cpp b/lib/Sema/CSTrail.cpp
diff --git a/lib/Sema/ConstraintGraph.cpp b/lib/Sema/ConstraintGraph.cpp
