Sema: Record expression patterns in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -40,6 +40,7 @@ LOCATOR_CHANGE(ResolvedOverload)
 
 EXPR_CHANGE(AppliedPropertyWrapper)
 EXPR_CHANGE(RecordedImpliedResult)
+EXPR_CHANGE(RecordedExprPattern)
 
 CHANGE(AddedTypeVariable)
 CHANGE(AddedConstraint)

include/swift/Sema/ConstraintSystem.h

@@ -1591,7 +1591,7 @@ class Solution {
 
   /// A map of expressions to the ExprPatterns that they are being solved as
   /// a part of.
-  llvm::MapVector<Expr *, ExprPattern *> exprPatterns;
+  llvm::DenseMap<Expr *, ExprPattern *> exprPatterns;
 
   /// The set of parameters that have been inferred to be 'isolated'.
   std::vector<ParamDecl *> isolatedParams;
@@ -2289,7 +2289,7 @@ class ConstraintSystem {
 
   /// A map of expressions to the ExprPatterns that they are being solved as
   /// a part of.
-  llvm::SmallMapVector<Expr *, ExprPattern *, 2> exprPatterns;
+  llvm::SmallDenseMap<Expr *, ExprPattern *, 2> exprPatterns;
 
   /// The set of parameters that have been inferred to be 'isolated'.
   llvm::SmallSetVector<ParamDecl *, 2> isolatedParams;
@@ -2855,9 +2855,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c exprPatterns.
-    unsigned numExprPatterns;
-
     /// The length of \c isolatedParams.
     unsigned numIsolatedParams;
 
@@ -3401,11 +3398,19 @@ class ConstraintSystem {
 
   /// Record a given ExprPattern as the parent of its sub-expression.
   void setExprPatternFor(Expr *E, ExprPattern *EP) {
-    assert(E);
-    assert(EP);
-    auto inserted = exprPatterns.insert({E, EP}).second;
-    assert(inserted && "Mapping already defined?");
-    (void)inserted;
+    ASSERT(E);
+    ASSERT(EP);
+    bool inserted = exprPatterns.insert({E, EP}).second;
+    ASSERT(inserted);
+
+    if (solverState)
+      recordChange(SolverTrail::Change::RecordedExprPattern(E));
+  }
+
+  /// Record a given ExprPattern as the parent of its sub-expression.
+  void removeExprPatternFor(Expr *E) {
+    bool erased = exprPatterns.erase(E);
+    ASSERT(erased);
   }
 
   std::optional<CaseLabelItemInfo>

lib/Sema/CSSolver.cpp

@@ -409,8 +409,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     isolatedParams.insert(param);
   }
 
-  for (auto &pair : solution.exprPatterns)
-    exprPatterns.insert(pair);
+  for (auto &pair : solution.exprPatterns) {
+    if (exprPatterns.count(pair.first) == 0)
+      setExprPatternFor(pair.first, pair.second);
+  }
 
   for (auto closure : solution.preconcurrencyClosures) {
     preconcurrencyClosures.insert(closure);
@@ -695,7 +697,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numExprPatterns = cs.exprPatterns.size();
   numIsolatedParams = cs.isolatedParams.size();
   numPreconcurrencyClosures = cs.preconcurrencyClosures.size();
   numImplicitValueConversions = cs.ImplicitValueConversions.size();
@@ -737,9 +738,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any ExprPattern mappings.
-  truncate(cs.exprPatterns, numExprPatterns);
-
   // Remove any isolated parameters.
   truncate(cs.isolatedParams, numIsolatedParams);
 

lib/Sema/CSTrail.cpp

@@ -489,6 +489,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedPotentialThrowSite:
     cs.removePotentialThrowSite(TheCatchNode);
     break;
+
+  case ChangeKind::RecordedExprPattern:
+    cs.removeExprPatternFor(TheExpr);
+    break;
   }
 }