Sema: Record ASTNode types in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -76,8 +76,10 @@ class SolverTrail {
     /// Recorded the mapping from a pack element expression to its parent
     /// pack expansion expression.
     RecordedPackEnvironment,
-    /// Record a defaulted constraint at a locator.
+    /// Recorded a defaulted constraint at a locator.
     RecordedDefaultedConstraint,
+    /// Recorded an assignment of a type to an AST node.
+    RecordedNodeType,
   };
 
   /// A change made to the constraint system.
@@ -142,6 +144,11 @@ class SolverTrail {
         Type ReqTy;
       } FixedRequirement;
 
+      struct {
+        ASTNode Node;
+        Type OldType;
+      } Node;
+
       ConstraintLocator *Locator;
       PackExpansionType *ExpansionTy;
       PackElementExpr *ElementExpr;
@@ -221,6 +228,9 @@ class SolverTrail {
     /// Create a change that recorded a defaulted constraint at a locator.
     static Change recordedDefaultedConstraint(ConstraintLocator *locator);
 
+    /// Create a change that recorded an assignment of a type to an AST node.
+    static Change recordedNodeType(ASTNode node, Type oldType);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2256,11 +2256,7 @@ class ConstraintSystem {
   /// nodes themselves. This allows us to typecheck and
   /// run through various diagnostics passes without actually mutating
   /// the types on the nodes.
-  llvm::MapVector<ASTNode, Type> NodeTypes;
-
-  /// The nodes for which we have produced types, along with the prior type
-  /// each node had before introducing this type.
-  llvm::SmallVector<std::pair<ASTNode, Type>, 8> addedNodeTypes;
+  llvm::DenseMap<ASTNode, Type> NodeTypes;
 
   /// Maps components in a key path expression to their type. Needed because
   /// KeyPathExpr + Index isn't an \c ASTNode and thus can't be stored in \c
@@ -2894,8 +2890,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    unsigned numAddedNodeTypes;
-
     unsigned numAddedKeyPathComponentTypes;
 
     unsigned numDisabledConstraints;
@@ -3205,21 +3199,39 @@ class ConstraintSystem {
     this->FavoredTypes[E] = T;
   }
 
-  /// Set the type in our type map for the given node.
+  /// Set the type in our type map for the given node, and record the change
+  /// in the trail.
   ///
-  /// The side tables are used through the expression type checker to avoid mutating nodes until
-  /// we know we have successfully type-checked them.
+  /// The side tables are used through the expression type checker to avoid
+  /// mutating nodes until we know we have successfully type-checked them.
   void setType(ASTNode node, Type type) {
-    assert(!node.isNull() && "Cannot set type information on null node");
-    assert(type && "Expected non-null type");
+    ASSERT(!node.isNull() && "Cannot set type information on null node");
+    ASSERT(type && "Expected non-null type");
 
     // Record the type.
     Type &entry = NodeTypes[node];
     Type oldType = entry;
     entry = type;
 
-    // Record the fact that we ascribed a type to this node.
-    addedNodeTypes.push_back({node, oldType});
+    if (oldType.getPointer() != type.getPointer()) {
+      // Record the fact that we assigned a type to this node.
+      if (isRecordingChanges())
+        recordChange(SolverTrail::Change::recordedNodeType(node, oldType));
+    }
+  }
+
+  /// Undo the above change.
+  void restoreType(ASTNode node, Type oldType) {
+    ASSERT(!node.isNull() && "Cannot set type information on null node");
+
+    if (oldType) {
+      auto found = NodeTypes.find(node);
+      ASSERT(found != NodeTypes.end());
+      found->second = oldType;
+    } else {
+      bool erased = NodeTypes.erase(node);
+      ASSERT(erased);
+    }
   }
 
   /// Set the type in our type map for a given expression. The side

lib/Sema/CSSolver.cpp

@@ -663,7 +663,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
 
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
-  numAddedNodeTypes = cs.addedNodeTypes.size();
   numAddedKeyPathComponentTypes = cs.addedKeyPathComponentTypes.size();
   numKeyPaths = cs.KeyPaths.size();
   numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
@@ -718,17 +717,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  // Remove any node types we registered.
-  for (unsigned i :
-           reverse(range(numAddedNodeTypes, cs.addedNodeTypes.size()))) {
-    auto node = cs.addedNodeTypes[i].first;
-    if (Type oldType = cs.addedNodeTypes[i].second)
-      cs.NodeTypes[node] = oldType;
-    else
-      cs.NodeTypes.erase(node);
-  }
-  truncate(cs.addedNodeTypes, numAddedNodeTypes);
-
   // Remove any node types we registered.
   for (unsigned i : reverse(range(numAddedKeyPathComponentTypes,
                                   cs.addedKeyPathComponentTypes.size()))) {

lib/Sema/CSTrail.cpp

@@ -225,6 +225,15 @@ SolverTrail::Change::recordedDefaultedConstraint(ConstraintLocator *locator) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedNodeType(ASTNode node, Type oldType) {
+  Change result;
+  result.Kind = ChangeKind::RecordedNodeType;
+  result.Node.Node = node;
+  result.Node.OldType = oldType;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -313,6 +322,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedDefaultedConstraint:
     cs.removeDefaultedConstraint(Locator);
     break;
+
+  case ChangeKind::RecordedNodeType:
+    cs.restoreType(Node.Node, Node.OldType);
+    break;
   }
 }
 
@@ -468,6 +481,7 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     break;
 
   case ChangeKind::RecordedPackEnvironment:
+    // FIXME: Print short form of PackExpansionExpr
     out << "(recorded pack environment)\n";
     break;
 
@@ -476,6 +490,17 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     Locator->dump(&cs.getASTContext().SourceMgr, out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedNodeType:
+    out << "(recorded node type at ";
+    Node.Node.getStartLoc().print(out, cs.getASTContext().SourceMgr);
+    out << " previous ";
+    if (Node.OldType)
+      Node.OldType->print(out, PO);
+    else
+      out << "null";
+    out << ")\n";
+    break;
   }
 }