Sema: Record argument matching choices in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -63,6 +63,8 @@ class SolverTrail {
     RecordedDisjunctionChoice,
     /// Recorded an applied disjunction.
     RecordedAppliedDisjunction,
+    /// Recorded an argument matching choice.
+    RecordedMatchCallArgumentResult,
   };
 
   /// A change made to the constraint system.
@@ -182,6 +184,9 @@ class SolverTrail {
     /// Create a change that recorded an applied disjunction.
     static Change recordedAppliedDisjunction(ConstraintLocator *locator);
 
+    /// Create a change that recorded an applied disjunction.
+    static Change recordedMatchCallArgumentResult(ConstraintLocator *locator);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -1519,7 +1519,7 @@ class Solution {
 
   /// For locators associated with call expressions, the trailing closure
   /// matching rule and parameter bindings that were applied.
-  llvm::MapVector<ConstraintLocator *, MatchCallArgumentResult>
+  llvm::DenseMap<ConstraintLocator *, MatchCallArgumentResult>
       argumentMatchingChoices;
 
   /// The set of disjunction choices used to arrive at this solution,
@@ -2341,7 +2341,7 @@ class ConstraintSystem {
 
   /// For locators associated with call expressions, the trailing closure
   /// matching rule and parameter bindings that were applied.
-  llvm::MapVector<ConstraintLocator *, MatchCallArgumentResult>
+  llvm::DenseMap<ConstraintLocator *, MatchCallArgumentResult>
       argumentMatchingChoices;
 
   /// The set of implicit value conversions performed by the solver on
@@ -2883,9 +2883,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c argumentMatchingChoices.
-    unsigned numArgumentMatchingChoices;
-
     /// The length of \c OpenedTypes.
     unsigned numOpenedTypes;
 
@@ -3637,9 +3634,16 @@ class ConstraintSystem {
   /// \param type The type on which to holeify.
   void recordTypeVariablesAsHoles(Type type);
 
+  /// Add a MatchCallArgumentResult and record the change in the trail.
   void recordMatchCallArgumentResult(ConstraintLocator *locator,
                                      MatchCallArgumentResult result);
 
+  /// Undo the above change.
+  void removeMatchCallArgumentResult(ConstraintLocator *locator) {
+    bool erased = argumentMatchingChoices.erase(locator);
+    ASSERT(erased);
+  }
+
   /// Record implicitly generated `callAsFunction` with root at the
   /// given expression, located at \c locator.
   void recordCallAsFunction(UnresolvedDotExpr *root, ArgumentList *arguments,

lib/Sema/CSSimplify.cpp

@@ -14925,7 +14925,11 @@ void ConstraintSystem::recordTypeVariablesAsHoles(Type type) {
 void ConstraintSystem::recordMatchCallArgumentResult(
     ConstraintLocator *locator, MatchCallArgumentResult result) {
   assert(locator->isLastElement<LocatorPathElt::ApplyArgument>());
-  argumentMatchingChoices.insert({locator, result});
+  bool inserted = argumentMatchingChoices.insert({locator, result}).second;
+  if (inserted) {
+    if (isRecordingChanges())
+      recordChange(SolverTrail::Change::recordedMatchCallArgumentResult(locator));
+  }
 }
 
 void ConstraintSystem::recordCallAsFunction(UnresolvedDotExpr *root,

lib/Sema/CSSolver.cpp

@@ -317,7 +317,7 @@ void ConstraintSystem::applySolution(const Solution &solution) {
 
   // Remember all of the argument/parameter matching choices we made.
   for (auto &argumentMatch : solution.argumentMatchingChoices) {
-    argumentMatchingChoices.insert(argumentMatch);
+    recordMatchCallArgumentResult(argumentMatch.first, argumentMatch.second);
   }
 
   // Remember implied results.
@@ -672,7 +672,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
 
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
-  numArgumentMatchingChoices = cs.argumentMatchingChoices.size();
   numOpenedTypes = cs.OpenedTypes.size();
   numOpenedExistentialTypes = cs.OpenedExistentialTypes.size();
   numOpenedPackExpansionTypes = cs.OpenedPackExpansionTypes.size();
@@ -735,9 +734,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  // Remove any argument matching choices;
-  truncate(cs.argumentMatchingChoices, numArgumentMatchingChoices);
-
   // Remove any opened types.
   truncate(cs.OpenedTypes, numOpenedTypes);
 

lib/Sema/CSTrail.cpp

@@ -169,6 +169,14 @@ SolverTrail::Change::recordedAppliedDisjunction(ConstraintLocator *locator) {
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedMatchCallArgumentResult(ConstraintLocator *locator) {
+  Change result;
+  result.Kind = ChangeKind::RecordedMatchCallArgumentResult;
+  result.Locator = locator;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -229,6 +237,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedAppliedDisjunction:
     cs.removeAppliedDisjunction(Locator);
     break;
+
+  case ChangeKind::RecordedMatchCallArgumentResult:
+    cs.removeMatchCallArgumentResult(Locator);
+    break;
   }
 }
 
@@ -352,6 +364,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     Locator->dump(&cs.getASTContext().SourceMgr, out);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedMatchCallArgumentResult:
+    out << "(recorded argument matching choice at ";
+    Locator->dump(&cs.getASTContext().SourceMgr, out);
+    out << ")\n";
+    break;
   }
 }