Sema: Record closure types in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -54,8 +54,9 @@ CHANGE(DisabledConstraint)
 CHANGE(FavoredConstraint)
 CHANGE(RecordedResultBuilderTransform)
 CHANGE(AppliedPropertyWrapper)
+CHANGE(RecordedClosureType)
 
-LAST_CHANGE(AppliedPropertyWrapper)
+LAST_CHANGE(RecordedClosureType)
 
 #undef LOCATOR_CHANGE
 #undef LAST_CHANGE

include/swift/Sema/CSTrail.h

@@ -124,6 +124,7 @@ class SolverTrail {
       PackElementExpr *ElementExpr;
       Expr *TheExpr;
       AnyFunctionRef AFR;
+      const ClosureExpr *Closure;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -205,6 +206,9 @@ class SolverTrail {
     /// Create a change that recorded a result builder transform.
     static Change AppliedPropertyWrapper(Expr *anchor);
 
+    /// Create a change that recorded a closure type.
+    static Change RecordedClosureType(const ClosureExpr *closure);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2221,6 +2221,9 @@ class ConstraintSystem {
 
   /// Maps discovered closures to their types inferred
   /// from declared parameters/result and body.
+  ///
+  /// This is a MapVector because contractEdges() iterates over it and
+  /// may depend on order.
   llvm::MapVector<const ClosureExpr *, FunctionType *> ClosureTypes;
 
   /// Maps closures and local functions to the pack expansion expressions they
@@ -2857,9 +2860,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c ClosureTypes.
-    unsigned numInferredClosureTypes;
-
     /// The length of \c ImpliedResults.
     unsigned numImpliedResults;
 
@@ -3083,10 +3083,18 @@ class ConstraintSystem {
   }
 
   void setClosureType(const ClosureExpr *closure, FunctionType *type) {
-    assert(closure);
-    assert(type && "Expected non-null type");
-    assert(ClosureTypes.count(closure) == 0 && "Cannot reset closure type");
-    ClosureTypes.insert({closure, type});
+    ASSERT(closure);
+    ASSERT(type);
+    bool inserted = ClosureTypes.insert({closure, type}).second;
+    ASSERT(inserted);
+
+    if (solverState)
+      recordChange(SolverTrail::Change::RecordedClosureType(closure));
+  }
+
+  void removeClosureType(const ClosureExpr *closure) {
+    bool erased = ClosureTypes.erase(closure);
+    ASSERT(erased);
   }
 
   FunctionType *getClosureType(const ClosureExpr *closure) const {

lib/Sema/CSSolver.cpp

@@ -690,7 +690,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numInferredClosureTypes = cs.ClosureTypes.size();
   numImpliedResults = cs.ImpliedResults.size();
   numContextualTypes = cs.contextualTypes.size();
   numTargets = cs.targets.size();
@@ -738,9 +737,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any inferred closure types (e.g. used in result builder body).
-  truncate(cs.ClosureTypes, numInferredClosureTypes);
-
   // Remove any implied results.
   truncate(cs.ImpliedResults, numImpliedResults);
 

lib/Sema/CSTrail.cpp

@@ -247,6 +247,14 @@ SolverTrail::Change::AppliedPropertyWrapper(Expr *expr) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::RecordedClosureType(const ClosureExpr *closure) {
+  Change result;
+  result.Kind = ChangeKind::RecordedClosureType;
+  result.Closure = closure;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -364,6 +372,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::ResolvedOverload:
     cs.removeResolvedOverload(Locator);
     break;
+
+  case ChangeKind::RecordedClosureType:
+    cs.removeClosureType(Closure);
+    break;
   }
 }
 
@@ -551,6 +563,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     simple_display(out, TheExpr);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedClosureType:
+    out << "(RecordedClosureType ";
+    simple_display(out, Closure);
+    out << ")\n";
+    break;
   }
 }