Sema: Record preconcurrency closures in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -26,6 +26,10 @@
 #define EXPR_CHANGE(Name) CHANGE(Name)
 #endif
 
+#ifndef CLOSURE_CHANGE
+#define CLOSURE_CHANGE(Name) CHANGE(Name)
+#endif
+
 #ifndef LAST_CHANGE
 #define LAST_CHANGE(Name)
 #endif
@@ -42,6 +46,9 @@ EXPR_CHANGE(AppliedPropertyWrapper)
 EXPR_CHANGE(RecordedImpliedResult)
 EXPR_CHANGE(RecordedExprPattern)
 
+CLOSURE_CHANGE(RecordedClosureType)
+CLOSURE_CHANGE(RecordedPreconcurrencyClosure)
+
 CHANGE(AddedTypeVariable)
 CHANGE(AddedConstraint)
 CHANGE(RemovedConstraint)
@@ -61,7 +68,6 @@ CHANGE(RecordedKeyPathComponentType)
 CHANGE(DisabledConstraint)
 CHANGE(FavoredConstraint)
 CHANGE(RecordedResultBuilderTransform)
-CHANGE(RecordedClosureType)
 CHANGE(RecordedContextualInfo)
 CHANGE(RecordedTarget)
 CHANGE(RecordedCaseLabelItemInfo)
@@ -72,5 +78,6 @@ LAST_CHANGE(RecordedIsolatedParam)
 
 #undef LOCATOR_CHANGE
 #undef EXPR_CHANGE
+#undef CLOSURE_CHANGE
 #undef LAST_CHANGE
 #undef CHANGE

include/swift/Sema/CSTrail.h

@@ -141,6 +141,7 @@ class SolverTrail {
 
 #define LOCATOR_CHANGE(Name) static Change Name(ConstraintLocator *locator);
 #define EXPR_CHANGE(Name) static Change Name(Expr *expr);
+#define CLOSURE_CHANGE(Name) static Change Name(ClosureExpr *closure);
 #include "swift/Sema/CSTrail.def"
 
     /// Create a change that added a type variable.
@@ -214,9 +215,6 @@ class SolverTrail {
     /// Create a change that recorded a result builder transform.
     static Change RecordedResultBuilderTransform(AnyFunctionRef fn);
 
-    /// Create a change that recorded a closure type.
-    static Change RecordedClosureType(ClosureExpr *closure);
-
     /// Create a change that recorded the contextual type of an AST node.
     static Change RecordedContextualInfo(ASTNode node);
 

include/swift/Sema/ConstraintSystem.h

@@ -1598,7 +1598,7 @@ class Solution {
 
   /// The set of closures that have been inferred to be "isolated by
   /// preconcurrency".
-  std::vector<const ClosureExpr *> preconcurrencyClosures;
+  llvm::DenseSet<const ClosureExpr *> preconcurrencyClosures;
 
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
@@ -2296,7 +2296,7 @@ class ConstraintSystem {
 
   /// The set of closures that have been inferred to be "isolated by
   /// preconcurrency".
-  llvm::SmallSetVector<const ClosureExpr *, 2> preconcurrencyClosures;
+  llvm::SmallDenseSet<const ClosureExpr *, 2> preconcurrencyClosures;
 
   /// Maps closure parameters to type variables.
   llvm::DenseMap<const ParamDecl *, TypeVariableType *>
@@ -2855,9 +2855,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c PreconcurrencyClosures.
-    unsigned numPreconcurrencyClosures;
-
     /// The length of \c ImplicitValueConversions.
     unsigned numImplicitValueConversions;
 
@@ -4219,6 +4216,13 @@ class ConstraintSystem {
   /// Undo the above change.
   void removeIsolatedParam(ParamDecl *param);
 
+  /// Used by the above to update preconcurrencyClosures and record a change in
+  /// the trail.
+  void recordPreconcurrencyClosure(const ClosureExpr *closure);
+
+  /// Undo the above change.
+  void removePreconcurrencyClosure(const ClosureExpr *closure);
+
   /// Given the fact that contextual type is now available for the type
   /// variable representing a pack expansion type, let's resolve the expansion.
   ///

lib/Sema/CSSimplify.cpp

@@ -11526,6 +11526,23 @@ void ConstraintSystem::removeIsolatedParam(ParamDecl *param) {
   ASSERT(erased);
 }
 
+void ConstraintSystem::recordPreconcurrencyClosure(
+    const ClosureExpr *closure) {
+  bool inserted = preconcurrencyClosures.insert(closure).second;
+  ASSERT(inserted);
+
+  if (solverState) {
+    recordChange(SolverTrail::Change::RecordedPreconcurrencyClosure(
+      const_cast<ClosureExpr *>(closure)));
+  }
+}
+
+void ConstraintSystem::removePreconcurrencyClosure(
+    const ClosureExpr *closure) {
+  bool erased = preconcurrencyClosures.erase(closure);
+  ASSERT(erased);
+}
+
 bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
                                       Type contextualType,
                                       ConstraintLocatorBuilder locator) {
@@ -11535,7 +11552,7 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
 
   // Note if this closure is isolated by preconcurrency.
   if (hasPreconcurrencyCallee(locator))
-    preconcurrencyClosures.insert(closure);
+    recordPreconcurrencyClosure(closure);
 
   // Let's look through all optionals associated with contextual
   // type to make it possible to infer parameter/result type of

lib/Sema/CSSolver.cpp

@@ -233,8 +233,8 @@ Solution ConstraintSystem::finalize() {
   for (const auto &param : isolatedParams)
     solution.isolatedParams.insert(param);
 
-  for (const auto &closure : preconcurrencyClosures)
-    solution.preconcurrencyClosures.push_back(closure);
+  for (auto closure : preconcurrencyClosures)
+    solution.preconcurrencyClosures.insert(closure);
 
   for (const auto &transformed : resultBuilderTransformed) {
     solution.resultBuilderTransformed.insert(transformed);
@@ -416,7 +416,8 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   }
 
   for (auto closure : solution.preconcurrencyClosures) {
-    preconcurrencyClosures.insert(closure);
+    if (preconcurrencyClosures.count(closure) == 0)
+      recordPreconcurrencyClosure(closure);
   }
 
   for (const auto &transformed : solution.resultBuilderTransformed) {
@@ -698,7 +699,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numPreconcurrencyClosures = cs.preconcurrencyClosures.size();
   numImplicitValueConversions = cs.ImplicitValueConversions.size();
   numArgumentLists = cs.ArgumentLists.size();
   numImplicitCallAsFunctionRoots = cs.ImplicitCallAsFunctionRoots.size();
@@ -738,9 +738,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any preconcurrency closures.
-  truncate(cs.preconcurrencyClosures, numPreconcurrencyClosures);
-
   // Remove any implicit value conversions.
   truncate(cs.ImplicitValueConversions, numImplicitValueConversions);
 

lib/Sema/CSSyntacticElement.cpp

@@ -2581,7 +2581,7 @@ bool ConstraintSystem::applySolution(AnyFunctionRef fn,
         param->setIsolated(true);
     }
 
-    if (llvm::is_contained(solution.preconcurrencyClosures, closure))
+    if (solution.preconcurrencyClosures.count(closure))
       closure->setIsolatedByPreconcurrency();
 
     // Coerce the result type, if it was written explicitly.

lib/Sema/CSTrail.cpp

@@ -64,6 +64,14 @@ SolverTrail::~SolverTrail() {
     result.TheExpr = expr; \
     return result; \
   }
+#define CLOSURE_CHANGE(Name) \
+  SolverTrail::Change \
+  SolverTrail::Change::Name(ClosureExpr *closure) { \
+    Change result; \
+    result.Kind = ChangeKind::Name; \
+    result.TheClosure = closure; \
+    return result; \
+  }
 #include "swift/Sema/CSTrail.def"
 
 SolverTrail::Change
@@ -247,14 +255,6 @@ SolverTrail::Change::RecordedResultBuilderTransform(AnyFunctionRef fn) {
   return result;
 }
 
-SolverTrail::Change
-SolverTrail::Change::RecordedClosureType(ClosureExpr *closure) {
-  Change result;
-  result.Kind = ChangeKind::RecordedClosureType;
-  result.TheClosure = closure;
-  return result;
-}
-
 SolverTrail::Change
 SolverTrail::Change::RecordedContextualInfo(ASTNode node) {
   Change result;
@@ -505,6 +505,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedIsolatedParam:
     cs.removeIsolatedParam(TheParam);
     break;
+
+  case ChangeKind::RecordedPreconcurrencyClosure:
+    cs.removePreconcurrencyClosure(TheClosure);
+    break;
   }
 }
 
@@ -530,6 +534,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     simple_display(out, TheExpr); \
     out << ")\n"; \
     break;
+#define CLOSURE_CHANGE(Name) \
+  case ChangeKind::Name: \
+    out << "(" << #Name << " "; \
+    simple_display(out, TheClosure); \
+    out << ")\n"; \
+    break;
 #include "swift/Sema/CSTrail.def"
 
   case ChangeKind::AddedTypeVariable:
@@ -693,12 +703,6 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     out << ")\n";
     break;
 
-  case ChangeKind::RecordedClosureType:
-    out << "(RecordedClosureType ";
-    simple_display(out, TheClosure);
-    out << ")\n";
-    break;
-
   case ChangeKind::RecordedContextualInfo:
     // FIXME: Print short form of ASTNode
     out << "(RecordedContextualInfo)\n";

lib/Sema/ConstraintSystem.cpp

@@ -4496,7 +4496,7 @@ size_t Solution::getTotalMemory() const {
          size_in_bytes(caseLabelItems) +
          size_in_bytes(exprPatterns) +
          size_in_bytes(isolatedParams) +
-         (preconcurrencyClosures.size() * sizeof(void *)) +
+         size_in_bytes(preconcurrencyClosures) +
          size_in_bytes(resultBuilderTransformed) +
          size_in_bytes(appliedPropertyWrappers) +
          size_in_bytes(argumentLists) +