Sema: Record key path expressions in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -74,8 +74,9 @@ CHANGE(RecordedTarget)
 CHANGE(RecordedCaseLabelItemInfo)
 CHANGE(RecordedPotentialThrowSite)
 CHANGE(RecordedIsolatedParam)
+CHANGE(RecordedKeyPath)
 
-LAST_CHANGE(RecordedIsolatedParam)
+LAST_CHANGE(RecordedKeyPath)
 
 #undef LOCATOR_CHANGE
 #undef EXPR_CHANGE

include/swift/Sema/CSTrail.h

@@ -226,6 +226,9 @@ class SolverTrail {
     /// Create a change that recorded an isolated parameter.
     static Change RecordedIsolatedParam(ParamDecl *param);
 
+    /// Create a change that recorded a key path expression.
+    static Change RecordedKeyPath(KeyPathExpr *expr);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -1567,9 +1567,9 @@ class Solution {
 
   /// The key path expression and its root type, value type, and decl context
   /// introduced by this solution.
-  llvm::MapVector<const KeyPathExpr *,
-                  std::tuple</*root=*/TypeVariableType *,
-                             /*value=*/TypeVariableType *, DeclContext *>>
+  llvm::DenseMap<const KeyPathExpr *,
+                 std::tuple</*root=*/TypeVariableType *,
+                            /*value=*/TypeVariableType *, DeclContext *>>
       KeyPaths;
 
   /// Contextual types introduced by this solution.
@@ -2265,9 +2265,9 @@ class ConstraintSystem {
       KeyPathComponentTypes;
 
   /// Maps a key path root, value, and decl context to the key path expression.
-  llvm::MapVector<const KeyPathExpr *,
-                  std::tuple</*root=*/TypeVariableType *,
-                             /*value=*/TypeVariableType *, DeclContext *>>
+  llvm::DenseMap<const KeyPathExpr *,
+                 std::tuple</*root=*/TypeVariableType *,
+                            /*value=*/TypeVariableType *, DeclContext *>>
       KeyPaths;
 
   /// Maps AST entries to their targets.
@@ -2850,9 +2850,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c KeyPaths.
-    unsigned numKeyPaths;
-
     /// The length of \c ArgumentLists.
     unsigned numArgumentLists;
 
@@ -3641,10 +3638,13 @@ class ConstraintSystem {
                             ConstraintLocator *locator);
 
   /// Record root, value, and declContext of keypath expression for use across
-  /// constraint system.
-  void recordKeyPath(KeyPathExpr *keypath, TypeVariableType *root,
+  /// constraint system, and add a change to the trail.
+  void recordKeyPath(const KeyPathExpr *keypath, TypeVariableType *root,
                      TypeVariableType *value, DeclContext *dc);
 
+  /// Undo the above change.
+  void removeKeyPath(const KeyPathExpr *keypath);
+
   /// Walk a closure AST to determine its effects.
   ///
   /// \returns a function's extended info describing the effects, as

lib/Sema/CSSimplify.cpp

@@ -14983,10 +14983,22 @@ void ConstraintSystem::recordCallAsFunction(UnresolvedDotExpr *root,
       getConstraintLocator(root, ConstraintLocator::ApplyArgument), arguments);
 }
 
-void ConstraintSystem::recordKeyPath(KeyPathExpr *keypath,
+void ConstraintSystem::recordKeyPath(const KeyPathExpr *keypath,
                                      TypeVariableType *root,
                                      TypeVariableType *value, DeclContext *dc) {
-  KeyPaths.insert(std::make_pair(keypath, std::make_tuple(root, value, dc)));
+  bool inserted = KeyPaths.insert(
+    std::make_pair(keypath, std::make_tuple(root, value, dc))).second;
+  ASSERT(inserted);
+
+  if (solverState) {
+    recordChange(SolverTrail::Change::RecordedKeyPath(
+      const_cast<KeyPathExpr *>(keypath)));
+  }
+}
+
+void ConstraintSystem::removeKeyPath(const KeyPathExpr *keypath) {
+  bool erased = KeyPaths.erase(keypath);
+  ASSERT(erased);
 }
 
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(

lib/Sema/CSSolver.cpp

@@ -378,7 +378,12 @@ void ConstraintSystem::applySolution(const Solution &solution) {
 
   // Add key paths.
   for (const auto &keypath : solution.KeyPaths) {
-    KeyPaths.insert(keypath);
+    if (KeyPaths.count(keypath.first) == 0) {
+      recordKeyPath(keypath.first,
+                    std::get<0>(keypath.second),
+                    std::get<1>(keypath.second),
+                    std::get<2>(keypath.second));
+    }
   }
 
   // Add the contextual types.
@@ -701,7 +706,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
 
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
-  numKeyPaths = cs.KeyPaths.size();
   numArgumentLists = cs.ArgumentLists.size();
   numImplicitCallAsFunctionRoots = cs.ImplicitCallAsFunctionRoots.size();
   numSynthesizedConformances = cs.SynthesizedConformances.size();
@@ -737,9 +741,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  /// Remove any key path expressions.
-  truncate(cs.KeyPaths, numKeyPaths);
-
   // Remove any argument lists no longer in scope.
   truncate(cs.ArgumentLists, numArgumentLists);
 

lib/Sema/CSTrail.cpp

@@ -317,6 +317,14 @@ SolverTrail::Change::RecordedIsolatedParam(ParamDecl *param) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::RecordedKeyPath(KeyPathExpr *expr) {
+  Change result;
+  result.Kind = ChangeKind::RecordedKeyPath;
+  result.KeyPath.Expr = expr;
+  return result;
+}
+
 SyntacticElementTargetKey
 SolverTrail::Change::getSyntacticElementTargetKey() const {
   ASSERT(Kind == ChangeKind::RecordedTarget);
@@ -475,6 +483,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedPreconcurrencyClosure:
     cs.removePreconcurrencyClosure(TheClosure);
     break;
+
+  case ChangeKind::RecordedKeyPath:
+    cs.removeKeyPath(KeyPath.Expr);
+    break;
   }
 }
 
@@ -688,6 +700,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     TheParam->dumpRef(out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedKeyPath:
+    out << "(RecordedKeyPath ";
+    simple_display(out, KeyPath.Expr);
+    out << ")\n";
+    break;
   }
 }