Sema: Record pack expansion environments in the trail

include/swift/Sema/CSTrail.h

@@ -71,6 +71,8 @@ public:
     RecordedOpenedExistentialType,
     /// Recorded the opening of a pack existential type.
     RecordedOpenedPackExpansionType,
+    /// Recorded the creation of a generic environment for a pack expansion expression.
+    RecordedPackExpansionEnvironment,
   };
 
   /// A change made to the constraint system.
@@ -203,6 +205,9 @@ public:
     /// Create a change that recorded the opening of a pack expansion type.
     static Change recordedOpenedPackExpansionType(PackExpansionType *expansion);
 
+    /// Create a change that recorded the opening of a pack expansion type.
+    static Change recordedPackExpansionEnvironment(ConstraintLocator *locator);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2397,7 +2397,7 @@ private:
   llvm::SmallDenseMap<PackExpansionType *, TypeVariableType *, 4>
       OpenedPackExpansionTypes;
 
-  llvm::SmallMapVector<ConstraintLocator *, std::pair<UUID, Type>, 4>
+  llvm::SmallDenseMap<ConstraintLocator *, std::pair<UUID, Type>, 4>
       PackExpansionEnvironments;
 
   llvm::SmallMapVector<PackElementExpr *, PackExpansionExpr *, 2>
@@ -2883,9 +2883,6 @@ public:
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c PackExpansionEnvironments.
-    unsigned numPackExpansionEnvironments;
-
     /// The length of \c PackEnvironments.
     unsigned numPackEnvironments;
 
@@ -3464,6 +3461,16 @@ public:
   GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator,
                                                 CanType shapeClass);
 
+  /// Update PackExpansionEnvironments and record a change in the trail.
+  void recordPackExpansionEnvironment(ConstraintLocator *locator,
+                                      std::pair<UUID, Type> uuidAndShape);
+
+  /// Undo the above change.
+  void removePackExpansionEnvironment(ConstraintLocator *locator) {
+    bool erased = PackExpansionEnvironments.erase(locator);
+    ASSERT(erased);
+  }
+
   /// Get the opened element generic environment for the given pack element.
   PackExpansionExpr *getPackEnvironment(PackElementExpr *packElement) const;
 

lib/Sema/CSSolver.cpp

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

lib/Sema/CSTrail.cpp

@@ -201,6 +201,14 @@ SolverTrail::Change::recordedOpenedPackExpansionType(PackExpansionType *expansio
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedPackExpansionEnvironment(ConstraintLocator *locator) {
+  Change result;
+  result.Kind = ChangeKind::RecordedPackExpansionEnvironment;
+  result.Locator = locator;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -277,6 +285,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedOpenedPackExpansionType:
     cs.removeOpenedPackExpansionType(ExpansionTy);
     break;
+
+  case ChangeKind::RecordedPackExpansionEnvironment:
+    cs.removePackExpansionEnvironment(Locator);
+    break;
   }
 }
 
@@ -424,6 +436,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     ExpansionTy->print(out, PO);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedPackExpansionEnvironment:
+    out << "(recorded pack expansion environment at ";
+    Locator->dump(&cs.getASTContext().SourceMgr, out);
+    out << ")\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -868,7 +868,7 @@ ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
   auto result = PackExpansionEnvironments.find(locator);
   if (result == PackExpansionEnvironments.end()) {
     uuidAndShape = std::make_pair(UUID::fromTime(), shapeClass);
-    PackExpansionEnvironments[locator] = uuidAndShape;
+    recordPackExpansionEnvironment(locator, uuidAndShape);
   } else {
     uuidAndShape = result->second;
   }
@@ -891,6 +891,17 @@ ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
                                               shapeParam, contextSubs);
 }
 
+void ConstraintSystem::recordPackExpansionEnvironment(
+    ConstraintLocator *locator, std::pair<UUID, Type> uuidAndShape) {
+  bool inserted = PackExpansionEnvironments.insert({locator, uuidAndShape}).second;
+  if (inserted) {
+    if (isRecordingChanges()) {
+      recordChange(
+        SolverTrail::Change::recordedPackExpansionEnvironment(locator));
+    }
+  }
+}
+
 PackExpansionExpr *
 ConstraintSystem::getPackEnvironment(PackElementExpr *packElement) const {
   const auto match = PackEnvironments.find(packElement);