Sema: Record opened pack expansion types in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -69,6 +69,8 @@ class SolverTrail {
     RecordedOpenedTypes,
     /// Recorded the opening of an existential type at a locator.
     RecordedOpenedExistentialType,
+    /// Recorded the opening of a pack existential type.
+    RecordedOpenedPackExpansionType,
   };
 
   /// A change made to the constraint system.
@@ -134,6 +136,7 @@ class SolverTrail {
       } FixedRequirement;
 
       ConstraintLocator *Locator;
+      PackExpansionType *ExpansionTy;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -197,6 +200,9 @@ class SolverTrail {
     /// Create a change that recorded the opening of an existential type.
     static Change recordedOpenedExistentialType(ConstraintLocator *locator);
 
+    /// Create a change that recorded the opening of a pack expansion type.
+    static Change recordedOpenedPackExpansionType(PackExpansionType *expansion);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2394,7 +2394,7 @@ class ConstraintSystem {
   llvm::SmallDenseMap<ConstraintLocator *, OpenedArchetypeType *, 4>
       OpenedExistentialTypes;
 
-  llvm::SmallMapVector<PackExpansionType *, TypeVariableType *, 4>
+  llvm::SmallDenseMap<PackExpansionType *, TypeVariableType *, 4>
       OpenedPackExpansionTypes;
 
   llvm::SmallMapVector<ConstraintLocator *, std::pair<UUID, Type>, 4>
@@ -2883,9 +2883,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c OpenedPackExpansionsTypes.
-    unsigned numOpenedPackExpansionTypes;
-
     /// The length of \c PackExpansionEnvironments.
     unsigned numPackExpansionEnvironments;
 
@@ -4335,6 +4332,16 @@ class ConstraintSystem {
                              OpenedTypeMap &replacements,
                              ConstraintLocatorBuilder locator);
 
+  /// Update OpenedPackExpansionTypes and record a change in the trail.
+  void recordOpenedPackExpansionType(PackExpansionType *expansion,
+                                     TypeVariableType *expansionVar);
+
+  /// Undo the above change.
+  void removeOpenedPackExpansionType(PackExpansionType *expansion) {
+    bool erased = OpenedPackExpansionTypes.erase(expansion);
+    ASSERT(erased);
+  }
+
 public:
   /// Recurse over the given type and open any opaque archetype types.
   Type openOpaqueType(Type type, ContextualTypePurpose context,

lib/Sema/CSSolver.cpp

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

lib/Sema/CSTrail.cpp

@@ -193,6 +193,14 @@ SolverTrail::Change::recordedOpenedExistentialType(ConstraintLocator *locator) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedOpenedPackExpansionType(PackExpansionType *expansionTy) {
+  Change result;
+  result.Kind = ChangeKind::RecordedOpenedPackExpansionType;
+  result.ExpansionTy = expansionTy;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -265,6 +273,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedOpenedExistentialType:
     cs.removeOpenedExistentialType(Locator);
     break;
+
+  case ChangeKind::RecordedOpenedPackExpansionType:
+    cs.removeOpenedPackExpansionType(ExpansionTy);
+    break;
   }
 }
 
@@ -406,6 +418,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     Locator->dump(&cs.getASTContext().SourceMgr, out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedOpenedPackExpansionType:
+    out << "(recorded opened pack expansion type for ";
+    ExpansionTy->print(out, PO);
+    out << ")\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -1263,6 +1263,15 @@ Type ConstraintSystem::openPackExpansionType(PackExpansionType *expansion,
   return expansionVar;
 }
 
+void ConstraintSystem::recordOpenedPackExpansionType(PackExpansionType *expansion,
+                                                     TypeVariableType *expansionVar) {
+  bool inserted = OpenedPackExpansionTypes.insert({expansion, expansionVar}).second;
+  if (inserted) {
+    if (isRecordingChanges())
+      recordChange(SolverTrail::Change::recordedOpenedPackExpansionType(expansion));
+  }
+}
+
 Type ConstraintSystem::openOpaqueType(OpaqueTypeArchetypeType *opaque,
                                       ConstraintLocatorBuilder locator) {
   auto opaqueDecl = opaque->getDecl();