Sema: Record opened types in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -65,6 +65,8 @@ class SolverTrail {
     RecordedAppliedDisjunction,
     /// Recorded an argument matching choice.
     RecordedMatchCallArgumentResult,
+    /// Recorded a list of opened types at a locator.
+    RecordedOpenedTypes,
   };
 
   /// A change made to the constraint system.
@@ -187,6 +189,9 @@ class SolverTrail {
     /// Create a change that recorded an applied disjunction.
     static Change recordedMatchCallArgumentResult(ConstraintLocator *locator);
 
+    /// Create a change that recorded a list of opened types.
+    static Change recordedOpenedTypes(ConstraintLocator *locator);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2362,7 +2362,7 @@ class ConstraintSystem {
 
   /// A mapping from constraint locators to the set of opened types associated
   /// with that locator.
-  llvm::SmallMapVector<ConstraintLocator *, ArrayRef<OpenedType>, 4>
+  llvm::SmallDenseMap<ConstraintLocator *, ArrayRef<OpenedType>, 4>
       OpenedTypes;
 
   /// A dictionary of all conformances that have been looked up by the solver.
@@ -2883,9 +2883,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c OpenedTypes.
-    unsigned numOpenedTypes;
-
     /// The length of \c OpenedExistentialTypes.
     unsigned numOpenedExistentialTypes;
 
@@ -4393,6 +4390,13 @@ class ConstraintSystem {
                               ConstraintLocatorBuilder locator,
                               llvm::function_ref<Type(Type)> subst);
 
+  /// Update OpenedTypes and record a change in the trail.
+  void recordOpenedType(
+      ConstraintLocator *locator, ArrayRef<OpenedType> openedTypes);
+
+  /// Undo the above change.
+  void removeOpenedType(ConstraintLocator *locator);
+
   /// Record the set of opened types for the given locator.
   void recordOpenedTypes(
          ConstraintLocatorBuilder locator,

lib/Sema/CSSolver.cpp

@@ -326,7 +326,7 @@ void ConstraintSystem::applySolution(const Solution &solution) {
 
   // Register the solution's opened types.
   for (const auto &opened : solution.OpenedTypes) {
-    OpenedTypes.insert(opened);
+    recordOpenedType(opened.first, opened.second);
   }
 
   // Register the solution's opened existential types.
@@ -672,7 +672,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
 
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
-  numOpenedTypes = cs.OpenedTypes.size();
   numOpenedExistentialTypes = cs.OpenedExistentialTypes.size();
   numOpenedPackExpansionTypes = cs.OpenedPackExpansionTypes.size();
   numPackExpansionEnvironments = cs.PackExpansionEnvironments.size();
@@ -734,9 +733,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  // Remove any opened types.
-  truncate(cs.OpenedTypes, numOpenedTypes);
-
   // Remove any opened existential types.
   truncate(cs.OpenedExistentialTypes, numOpenedExistentialTypes);
 

lib/Sema/CSTrail.cpp

@@ -177,6 +177,14 @@ SolverTrail::Change::recordedMatchCallArgumentResult(ConstraintLocator *locator)
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedOpenedTypes(ConstraintLocator *locator) {
+  Change result;
+  result.Kind = ChangeKind::RecordedOpenedTypes;
+  result.Locator = locator;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -241,6 +249,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedMatchCallArgumentResult:
     cs.removeMatchCallArgumentResult(Locator);
     break;
+
+  case ChangeKind::RecordedOpenedTypes:
+    cs.removeOpenedType(Locator);
+    break;
   }
 }
 
@@ -370,6 +382,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     Locator->dump(&cs.getASTContext().SourceMgr, out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedOpenedTypes:
+    out << "(recorded list of opened types at ";
+    Locator->dump(&cs.getASTContext().SourceMgr, out);
+    out << ")\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -1654,6 +1654,20 @@ Type ConstraintSystem::getUnopenedTypeOfReference(
   return requestedType;
 }
 
+void ConstraintSystem::recordOpenedType(
+    ConstraintLocator *locator, ArrayRef<OpenedType> openedTypes) {
+  bool inserted = OpenedTypes.insert({locator, openedTypes}).second;
+  if (inserted) {
+    if (isRecordingChanges())
+      recordChange(SolverTrail::Change::recordedOpenedTypes(locator));
+  }
+}
+
+void ConstraintSystem::removeOpenedType(ConstraintLocator *locator) {
+  bool erased = OpenedTypes.erase(locator);
+  ASSERT(erased);
+}
+
 void ConstraintSystem::recordOpenedTypes(
        ConstraintLocatorBuilder locator,
        const OpenedTypeMap &replacements) {
@@ -1673,20 +1687,12 @@ void ConstraintSystem::recordOpenedTypes(
 
   ConstraintLocator *locatorPtr = getConstraintLocator(locator);
   assert(locatorPtr && "No locator for opened types?");
-#if false
-  assert(std::find_if(OpenedTypes.begin(), OpenedTypes.end(),
-                      [&](const std::pair<ConstraintLocator *,
-                          ArrayRef<OpenedType>> &entry) {
-                        return entry.first == locatorPtr;
-                      }) == OpenedTypes.end() &&
-         "already registered opened types for this locator");
-#endif
 
   OpenedType* openedTypes
     = Allocator.Allocate<OpenedType>(replacements.size());
   std::copy(replacements.begin(), replacements.end(), openedTypes);
-  OpenedTypes.insert(
-      {locatorPtr, llvm::ArrayRef(openedTypes, replacements.size())});
+  recordOpenedType(
+      locatorPtr, llvm::ArrayRef(openedTypes, replacements.size()));
 }
 
 /// Determine how many levels of argument labels should be removed from the