Sema: Record key path component types in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -80,6 +80,8 @@ class SolverTrail {
     RecordedDefaultedConstraint,
     /// Recorded an assignment of a type to an AST node.
     RecordedNodeType,
+    /// Recorded an assignment of a type to a keypath component.
+    RecordedKeyPathComponentType,
   };
 
   /// A change made to the constraint system.
@@ -149,6 +151,11 @@ class SolverTrail {
         Type OldType;
       } Node;
 
+      struct {
+        const KeyPathExpr *Expr;
+        Type OldType;
+      } KeyPath;
+
       ConstraintLocator *Locator;
       PackExpansionType *ExpansionTy;
       PackElementExpr *ElementExpr;
@@ -231,6 +238,11 @@ class SolverTrail {
     /// Create a change that recorded an assignment of a type to an AST node.
     static Change recordedNodeType(ASTNode node, Type oldType);
 
+    /// Create a change that recorded an assignment of a type to an AST node.
+    static Change recordedKeyPathComponentType(const KeyPathExpr *expr,
+                                               unsigned component,
+                                               Type oldType);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2265,11 +2265,6 @@ class ConstraintSystem {
                  Type>
       KeyPathComponentTypes;
 
-  /// Same as \c addedNodeTypes for \c KeyPathComponentTypes.
-  llvm::SmallVector<
-      std::tuple<const KeyPathExpr *, /*component index=*/unsigned, Type>>
-      addedKeyPathComponentTypes;
-
   /// Maps a key path root, value, and decl context to the key path expression.
   llvm::MapVector<const KeyPathExpr *,
                   std::tuple</*root=*/TypeVariableType *,
@@ -2890,8 +2885,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    unsigned numAddedKeyPathComponentTypes;
-
     unsigned numDisabledConstraints;
 
     unsigned numFavoredConstraints;
@@ -3234,25 +3227,44 @@ class ConstraintSystem {
     }
   }
 
+  /// Check to see if we have a type for a node.
+  bool hasType(ASTNode node) const {
+    assert(!node.isNull() && "Expected non-null node");
+    return NodeTypes.count(node) > 0;
+  }
+
   /// Set the type in our type map for a given expression. The side
   /// map is used throughout the expression type checker in order to
   /// avoid mutating expressions until we know we have successfully
   /// type-checked them.
   void setType(const KeyPathExpr *KP, unsigned I, Type T) {
-    assert(KP && "Expected non-null key path parameter!");
-    assert(T && "Expected non-null type!");
+    ASSERT(KP && "Expected non-null key path parameter!");
+    ASSERT(T && "Expected non-null type!");
 
     Type &entry = KeyPathComponentTypes[{KP, I}];
     Type oldType = entry;
     entry = T;
 
-    addedKeyPathComponentTypes.push_back(std::make_tuple(KP, I, oldType));
+    if (oldType.getPointer() != T.getPointer()) {
+      if (isRecordingChanges()) {
+        recordChange(
+          SolverTrail::Change::recordedKeyPathComponentType(
+            KP, I, oldType));
+      }
+    }
   }
 
-  /// Check to see if we have a type for a node.
-  bool hasType(ASTNode node) const {
-    assert(!node.isNull() && "Expected non-null node");
-    return NodeTypes.count(node) > 0;
+  void restoreType(const KeyPathExpr *KP, unsigned I, Type T) {
+    ASSERT(KP && "Expected non-null key path parameter!");
+
+    if (T) {
+      auto found = KeyPathComponentTypes.find({KP, I});
+      ASSERT(found != KeyPathComponentTypes.end());
+      found->second = T;
+    } else {
+      bool erased = KeyPathComponentTypes.erase({KP, I});
+      ASSERT(erased);
+    }
   }
 
   bool hasType(const KeyPathExpr *KP, unsigned I) const {

lib/Sema/CSSolver.cpp

@@ -663,7 +663,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
 
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
-  numAddedKeyPathComponentTypes = cs.addedKeyPathComponentTypes.size();
   numKeyPaths = cs.KeyPaths.size();
   numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
   numFavoredConstraints = cs.solverState->getNumFavoredConstraints();
@@ -717,19 +716,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(numAddedKeyPathComponentTypes,
-                                  cs.addedKeyPathComponentTypes.size()))) {
-    auto KeyPath = std::get<0>(cs.addedKeyPathComponentTypes[i]);
-    auto KeyPathIndex = std::get<1>(cs.addedKeyPathComponentTypes[i]);
-    if (Type oldType = std::get<2>(cs.addedKeyPathComponentTypes[i])) {
-      cs.KeyPathComponentTypes[{KeyPath, KeyPathIndex}] = oldType;
-    } else {
-      cs.KeyPathComponentTypes.erase({KeyPath, KeyPathIndex});
-    }
-  }
-  truncate(cs.addedKeyPathComponentTypes, numAddedKeyPathComponentTypes);
-
   /// Remove any key path expressions.
   truncate(cs.KeyPaths, numKeyPaths);
 

lib/Sema/CSTrail.cpp

@@ -234,6 +234,18 @@ SolverTrail::Change::recordedNodeType(ASTNode node, Type oldType) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedKeyPathComponentType(const KeyPathExpr *expr,
+                                                  unsigned component,
+                                                  Type oldType) {
+  Change result;
+  result.Kind = ChangeKind::RecordedKeyPathComponentType;
+  result.Options = component;
+  result.KeyPath.Expr = expr;
+  result.KeyPath.OldType = oldType;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -326,6 +338,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedNodeType:
     cs.restoreType(Node.Node, Node.OldType);
     break;
+
+  case ChangeKind::RecordedKeyPathComponentType:
+    cs.restoreType(KeyPath.Expr, Options, KeyPath.OldType);
+    break;
   }
 }
 
@@ -482,7 +498,9 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
 
   case ChangeKind::RecordedPackEnvironment:
     // FIXME: Print short form of PackExpansionExpr
-    out << "(recorded pack environment)\n";
+    out << "(recorded pack environment";
+    simple_display(out, ElementExpr);
+    out << "\n";
     break;
 
   case ChangeKind::RecordedDefaultedConstraint:
@@ -501,6 +519,17 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
       out << "null";
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedKeyPathComponentType:
+    out << "(recorded key path ";
+    simple_display(out, KeyPath.Expr);
+    out << " with component type ";
+    if (Node.OldType)
+      Node.OldType->print(out, PO);
+    else
+      out << "null";
+    out << " for component " << Options << ")\n";
+    break;
   }
 }