Sema: Record pack environments in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -73,6 +73,9 @@ class SolverTrail {
     RecordedOpenedPackExpansionType,
     /// Recorded the creation of a generic environment for a pack expansion expression.
     RecordedPackExpansionEnvironment,
+    /// Recorded the mapping from a pack element expression to its parent
+    /// pack expansion expression.
+    RecordedPackEnvironment,
   };
 
   /// A change made to the constraint system.
@@ -139,6 +142,7 @@ class SolverTrail {
 
       ConstraintLocator *Locator;
       PackExpansionType *ExpansionTy;
+      PackElementExpr *ElementExpr;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -208,6 +212,10 @@ class SolverTrail {
     /// Create a change that recorded the opening of a pack expansion type.
     static Change recordedPackExpansionEnvironment(ConstraintLocator *locator);
 
+    /// Create a change that recorded a mapping from a pack element expression
+    /// to its parent expansion expression.
+    static Change recordedPackEnvironment(PackElementExpr *packElement);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -1546,7 +1546,7 @@ class Solution {
       PackExpansionEnvironments;
 
   /// The pack expansion environment that can open a given pack element.
-  llvm::MapVector<PackElementExpr *, PackExpansionExpr *>
+  llvm::DenseMap<PackElementExpr *, PackExpansionExpr *>
       PackEnvironments;
 
   /// The outer pack element generic environment to use when dealing with nested
@@ -2400,7 +2400,7 @@ class ConstraintSystem {
   llvm::SmallDenseMap<ConstraintLocator *, std::pair<UUID, Type>, 4>
       PackExpansionEnvironments;
 
-  llvm::SmallMapVector<PackElementExpr *, PackExpansionExpr *, 2>
+  llvm::SmallDenseMap<PackElementExpr *, PackExpansionExpr *, 2>
       PackEnvironments;
 
   llvm::SmallVector<GenericEnvironment *, 4> PackElementGenericEnvironments;
@@ -2883,9 +2883,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c PackEnvironments.
-    unsigned numPackEnvironments;
-
     /// The length of \c PackElementGenericEnvironments.
     unsigned numPackElementGenericEnvironments;
 
@@ -3474,10 +3471,17 @@ class ConstraintSystem {
   /// Get the opened element generic environment for the given pack element.
   PackExpansionExpr *getPackEnvironment(PackElementExpr *packElement) const;
 
-  /// Associate an opened element generic environment to a pack element.
+  /// Associate an opened element generic environment to a pack element,
+  /// and record a change in the trail.
   void addPackEnvironment(PackElementExpr *packElement,
                           PackExpansionExpr *packExpansion);
 
+  /// Undo the above change.
+  void removePackEnvironment(PackElementExpr *packElement) {
+    bool erased = PackEnvironments.erase(packElement);
+    ASSERT(erased);
+  }
+
   /// Retrieve the constraint locator for the given anchor and
   /// path, uniqued and automatically infer the summary flags
   ConstraintLocator *

lib/Sema/CSSolver.cpp

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

lib/Sema/CSTrail.cpp

@@ -209,6 +209,14 @@ SolverTrail::Change::recordedPackExpansionEnvironment(ConstraintLocator *locator
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedPackEnvironment(PackElementExpr *packElement) {
+  Change result;
+  result.Kind = ChangeKind::RecordedPackEnvironment;
+  result.ElementExpr = packElement;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -289,6 +297,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedPackExpansionEnvironment:
     cs.removePackExpansionEnvironment(Locator);
     break;
+
+  case ChangeKind::RecordedPackEnvironment:
+    cs.removePackEnvironment(ElementExpr);
+    break;
   }
 }
 
@@ -442,6 +454,10 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     Locator->dump(&cs.getASTContext().SourceMgr, out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedPackEnvironment:
+    out << "(recorded pack environment)\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -910,11 +910,12 @@ ConstraintSystem::getPackEnvironment(PackElementExpr *packElement) const {
 
 void ConstraintSystem::addPackEnvironment(PackElementExpr *packElement,
                                           PackExpansionExpr *packExpansion) {
-  assert(packElement);
-  assert(packExpansion);
-  [[maybe_unused]] const auto inserted =
+  bool inserted =
       PackEnvironments.insert({packElement, packExpansion}).second;
-  assert(inserted && "Mapping already defined?");
+  if (inserted) {
+    if (isRecordingChanges())
+      recordChange(SolverTrail::Change::recordedPackEnvironment(packElement));
+  }
 }
 
 /// Extend the given depth map by adding depths for all of the subexpressions