Sema: Record resolved overloads in the trail

Author: slavapestov

include/swift/Sema/CSTrail.def

@@ -32,6 +32,7 @@ LOCATOR_CHANGE(RecordedOpenedTypes)
 LOCATOR_CHANGE(RecordedOpenedExistentialType)
 LOCATOR_CHANGE(RecordedPackExpansionEnvironment)
 LOCATOR_CHANGE(RecordedDefaultedConstraint)
+LOCATOR_CHANGE(ResolvedOverload)
 
 CHANGE(AddedTypeVariable)
 CHANGE(AddedConstraint)

include/swift/Sema/ConstraintSystem.h

@@ -2207,7 +2207,7 @@ class ConstraintSystem {
   llvm::FoldingSetVector<ConstraintLocator> ConstraintLocators;
 
   /// The overload sets that have been resolved along the current path.
-  llvm::MapVector<ConstraintLocator *, SelectedOverload> ResolvedOverloads;
+  llvm::DenseMap<ConstraintLocator *, SelectedOverload> ResolvedOverloads;
 
   /// The current fixed score for this constraint system and the (partial)
   /// solution it represents.
@@ -2857,9 +2857,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c ResolvedOverloads.
-    unsigned numResolvedOverloads;
-
     /// The length of \c ClosureTypes.
     unsigned numInferredClosureTypes;
 
@@ -4932,6 +4929,11 @@ class ConstraintSystem {
     buildDisjunctionForOptionalVsUnderlying(boundTy, type, dynamicLocator);
   }
 
+  void recordResolvedOverload(ConstraintLocator *locator,
+                              SelectedOverload choice);
+
+  void removeResolvedOverload(ConstraintLocator *locator);
+
   /// Resolve the given overload set to the given choice.
   void resolveOverload(ConstraintLocator *locator, Type boundType,
                        OverloadChoice choice, DeclContext *useDC);
@@ -5689,7 +5691,7 @@ class ConstraintSystem {
   }
 
   /// The overload sets that have already been resolved along the current path.
-  const llvm::MapVector<ConstraintLocator *, SelectedOverload> &
+  const llvm::DenseMap<ConstraintLocator *, SelectedOverload> &
   getResolvedOverloads() const {
     return ResolvedOverloads;
   }

lib/Sema/CSSolver.cpp

@@ -290,8 +290,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
 
   // Register overload choices.
   // FIXME: Copy these directly into some kind of partial solution?
-  for (auto overload : solution.overloadChoices)
-    ResolvedOverloads.insert(overload);
+  for (auto overload : solution.overloadChoices) {
+    if (!ResolvedOverloads.count(overload.first))
+      recordResolvedOverload(overload.first, overload.second);
+  }
 
   // Register constraint restrictions.
   // FIXME: Copy these directly into some kind of partial solution?
@@ -688,7 +690,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
   numKeyPaths = cs.KeyPaths.size();
-  numResolvedOverloads = cs.ResolvedOverloads.size();
   numInferredClosureTypes = cs.ClosureTypes.size();
   numImpliedResults = cs.ImpliedResults.size();
   numContextualTypes = cs.contextualTypes.size();
@@ -717,8 +718,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Erase the end of various lists.
   truncate(cs.TypeVariables, numTypeVariables);
 
-  truncate(cs.ResolvedOverloads, numResolvedOverloads);
-
   // Move any remaining active constraints into the inactive list.
   if (!cs.ActiveConstraints.empty()) {
     for (auto &constraint : cs.ActiveConstraints) {

lib/Sema/CSTrail.cpp

@@ -360,6 +360,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::AppliedPropertyWrapper:
     cs.removePropertyWrapper(TheExpr);
     break;
+
+  case ChangeKind::ResolvedOverload:
+    cs.removeResolvedOverload(Locator);
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -3735,6 +3735,20 @@ void ConstraintSystem::bindOverloadType(
   llvm_unreachable("Unhandled OverloadChoiceKind in switch.");
 }
 
+void ConstraintSystem::recordResolvedOverload(ConstraintLocator *locator,
+                                              SelectedOverload overload) {
+  bool inserted = ResolvedOverloads.insert({locator, overload}).second;
+  ASSERT(inserted);
+
+  if (solverState)
+    recordChange(SolverTrail::Change::ResolvedOverload(locator));
+}
+
+void ConstraintSystem::removeResolvedOverload(ConstraintLocator *locator) {
+  bool erased = ResolvedOverloads.erase(locator);
+  ASSERT(erased);
+}
+
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
                                        Type boundType,
                                        OverloadChoice choice,
@@ -3990,9 +4004,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
   auto overload = SelectedOverload{
       choice, openedType, adjustedOpenedType, refType, adjustedRefType,
       boundType};
-  auto result = ResolvedOverloads.insert({locator, overload});
-  assert(result.second && "Already resolved this overload?");
-  (void)result;
+  recordResolvedOverload(locator, overload);
 
   // Add the constraints necessary to bind the overload type.
   bindOverloadType(overload, boundType, locator, useDC,