Sema: Record contextual types in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -61,8 +61,9 @@ CHANGE(DisabledConstraint)
 CHANGE(FavoredConstraint)
 CHANGE(RecordedResultBuilderTransform)
 CHANGE(RecordedClosureType)
+CHANGE(RecordedContextualInfo)
 
-LAST_CHANGE(RecordedClosureType)
+LAST_CHANGE(RecordedContextualInfo)
 
 #undef LOCATOR_CHANGE
 #undef EXPR_CHANGE

include/swift/Sema/CSTrail.h

@@ -207,6 +207,9 @@ class SolverTrail {
     /// Create a change that recorded a closure type.
     static Change RecordedClosureType(const ClosureExpr *closure);
 
+    /// Create a change that recorded the contextual type of an AST node.
+    static Change RecordedContextualInfo(ASTNode node);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2281,7 +2281,7 @@ class ConstraintSystem {
   /// constraint system. The second type, if valid, contains the type as it
   /// should appear in actual constraint. This will have unbound generic types
   /// opened, placeholder types converted to type variables, etc.
-  llvm::MapVector<ASTNode, std::pair<ContextualTypeInfo, Type>> contextualTypes;
+  llvm::DenseMap<ASTNode, std::pair<ContextualTypeInfo, Type>> contextualTypes;
 
   /// Information about each case label item tracked by the constraint system.
   llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
@@ -2860,9 +2860,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c contextualTypes.
-    unsigned numContextualTypes;
-
     /// The length of \c targets.
     unsigned numTargets;
 
@@ -3294,10 +3291,17 @@ class ConstraintSystem {
   }
 
   void setContextualInfo(ASTNode node, ContextualTypeInfo info) {
-    assert(bool(node) && "Expected non-null expression!");
-    assert(contextualTypes.count(node) == 0 &&
-           "Already set this contextual type");
-    contextualTypes[node] = {info, Type()};
+    ASSERT(bool(node) && "Expected non-null expression!");
+    bool inserted = contextualTypes.insert({node, {info, Type()}}).second;
+    ASSERT(inserted);
+
+    if (solverState)
+      recordChange(SolverTrail::Change::RecordedContextualInfo(node));
+  }
+
+  void removeContextualInfo(ASTNode node) {
+    bool erased = contextualTypes.erase(node);
+    ASSERT(erased);
   }
 
   std::optional<ContextualTypeInfo> getContextualTypeInfo(ASTNode node) const {

lib/Sema/CSSolver.cpp

@@ -692,7 +692,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numContextualTypes = cs.contextualTypes.size();
   numTargets = cs.targets.size();
   numCaseLabelItems = cs.caseLabelItems.size();
   numPotentialThrowSites = cs.potentialThrowSites.size();
@@ -738,9 +737,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 
-  // Remove any contextual types.
-  truncate(cs.contextualTypes, numContextualTypes);
-
   // Remove any targets.
   truncate(cs.targets, numTargets);
 

lib/Sema/CSTrail.cpp

@@ -255,6 +255,14 @@ SolverTrail::Change::RecordedClosureType(const ClosureExpr *closure) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::RecordedContextualInfo(ASTNode node) {
+  Change result;
+  result.Kind = ChangeKind::RecordedContextualInfo;
+  result.Node.Node = node;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -380,6 +388,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedImpliedResult:
     cs.removeImpliedResult(TheExpr);
     break;
+
+  case ChangeKind::RecordedContextualInfo:
+    cs.removeContextualInfo(Node.Node);
+    break;
   }
 }
 
@@ -573,6 +585,11 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     simple_display(out, Closure);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedContextualInfo:
+    // FIXME: Print short form of ASTNode
+    out << "(RecordedContextualInfo)\n";
+    break;
   }
 }