Sema: Record opened existential types in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -67,6 +67,8 @@ class SolverTrail {
     RecordedMatchCallArgumentResult,
     /// Recorded a list of opened types at a locator.
     RecordedOpenedTypes,
+    /// Recorded the opening of an existential type at a locator.
+    RecordedOpenedExistentialType,
   };
 
   /// A change made to the constraint system.
@@ -192,6 +194,9 @@ class SolverTrail {
     /// Create a change that recorded a list of opened types.
     static Change recordedOpenedTypes(ConstraintLocator *locator);
 
+    /// Create a change that recorded the opening of an existential type.
+    static Change recordedOpenedExistentialType(ConstraintLocator *locator);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2391,7 +2391,7 @@ class ConstraintSystem {
 
   /// A mapping from constraint locators to the opened existential archetype
   /// used for the 'self' of an existential type.
-  llvm::SmallMapVector<ConstraintLocator *, OpenedArchetypeType *, 4>
+  llvm::SmallDenseMap<ConstraintLocator *, OpenedArchetypeType *, 4>
       OpenedExistentialTypes;
 
   llvm::SmallMapVector<PackExpansionType *, TypeVariableType *, 4>
@@ -2883,9 +2883,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c OpenedExistentialTypes.
-    unsigned numOpenedExistentialTypes;
-
     /// The length of \c OpenedPackExpansionsTypes.
     unsigned numOpenedPackExpansionTypes;
 
@@ -3456,6 +3453,16 @@ class ConstraintSystem {
   std::pair<Type, OpenedArchetypeType *> openExistentialType(
       Type type, ConstraintLocator *locator);
 
+  /// Update OpenedExistentials and record a change in the trail.
+  void recordOpenedExistentialType(ConstraintLocator *locator,
+                                   OpenedArchetypeType *opened);
+
+  /// Undo the above change.
+  void removeOpenedExistentialType(ConstraintLocator *locator) {
+    bool erased = OpenedExistentialTypes.erase(locator);
+    ASSERT(erased);
+  }
+
   /// Get the opened element generic environment for the given locator.
   GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator,
                                                 CanType shapeClass);

lib/Sema/CSSolver.cpp

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

lib/Sema/CSTrail.cpp

@@ -185,6 +185,14 @@ SolverTrail::Change::recordedOpenedTypes(ConstraintLocator *locator) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedOpenedExistentialType(ConstraintLocator *locator) {
+  Change result;
+  result.Kind = ChangeKind::RecordedOpenedExistentialType;
+  result.Locator = locator;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -253,6 +261,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedOpenedTypes:
     cs.removeOpenedType(Locator);
     break;
+
+  case ChangeKind::RecordedOpenedExistentialType:
+    cs.removeOpenedExistentialType(Locator);
+    break;
   }
 }
 
@@ -388,6 +400,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     Locator->dump(&cs.getASTContext().SourceMgr, out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedOpenedExistentialType:
+    out << "(recorded opened existential type at ";
+    Locator->dump(&cs.getASTContext().SourceMgr, out);
+    out << ")\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -845,11 +845,20 @@ std::pair<Type, OpenedArchetypeType *> ConstraintSystem::openExistentialType(
     t = t->getMetatypeInstanceType();
   auto *opened = t->castTo<OpenedArchetypeType>();
 
-  assert(OpenedExistentialTypes.count(locator) == 0);
-  OpenedExistentialTypes.insert({locator, opened});
+  recordOpenedExistentialType(locator, opened);
+
   return {result, opened};
 }
 
+void ConstraintSystem::recordOpenedExistentialType(
+    ConstraintLocator *locator, OpenedArchetypeType *opened) {
+  bool inserted = OpenedExistentialTypes.insert({locator, opened}).second;
+  if (inserted) {
+    if (isRecordingChanges())
+      recordChange(SolverTrail::Change::recordedOpenedExistentialType(locator));
+  }
+}
+
 GenericEnvironment *
 ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
                                             CanType shapeClass) {
@@ -2740,8 +2749,7 @@ DeclReferenceType ConstraintSystem::getTypeOfMemberReference(
   } else if (baseObjTy->isExistentialType()) {
     auto openedArchetype =
         OpenedArchetypeType::get(baseObjTy->getCanonicalType());
-    OpenedExistentialTypes.insert(
-        {getConstraintLocator(locator), openedArchetype});
+    recordOpenedExistentialType(getConstraintLocator(locator), openedArchetype);
     baseOpenedTy = openedArchetype;
   }