Sema: Record disjunction choices in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -59,6 +59,8 @@ class SolverTrail {
     AddedFix,
     /// Recorded a fixed requirement.
     AddedFixedRequirement,
+    /// Recorded a disjunction choice.
+    RecordedDisjunctionChoice,
   };
 
   /// A change made to the constraint system.
@@ -122,6 +124,8 @@ class SolverTrail {
         GenericTypeParamType *GP;
         Type ReqTy;
       } FixedRequirement;
+
+      ConstraintLocator *Locator;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -169,6 +173,10 @@ class SolverTrail {
                                         unsigned reqKind,
                                         Type requirementTy);
 
+    /// Create a change that recorded a disjunction choice.
+    static Change recordedDisjunctionChoice(ConstraintLocator *locator,
+                                            unsigned index);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2328,7 +2328,7 @@ class ConstraintSystem {
 
   /// The set of remembered disjunction choices used to reach
   /// the current constraint system.
-  llvm::MapVector<ConstraintLocator *, unsigned> DisjunctionChoices;
+  llvm::SmallDenseMap<ConstraintLocator *, unsigned, 4> DisjunctionChoices;
 
   /// A map from applied disjunction constraints to the corresponding
   /// argument function type.
@@ -2879,9 +2879,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c DisjunctionChoices.
-    unsigned numDisjunctionChoices;
-
     /// The length of \c AppliedDisjunctions.
     unsigned numAppliedDisjunctions;
 
@@ -5319,15 +5316,16 @@ class ConstraintSystem {
   /// Collect the current inactive disjunction constraints.
   void collectDisjunctions(SmallVectorImpl<Constraint *> &disjunctions);
 
-  /// Record a particular disjunction choice of
-  void recordDisjunctionChoice(ConstraintLocator *disjunctionLocator,
-                               unsigned index) {
-    // We shouldn't ever register disjunction choices multiple times.
-    assert(!DisjunctionChoices.count(disjunctionLocator) ||
-           DisjunctionChoices[disjunctionLocator] == index);
-    DisjunctionChoices.insert({disjunctionLocator, index});
+  /// Record a particular disjunction choice and add a change to the trail.
+  void recordDisjunctionChoice(ConstraintLocator *locator, unsigned index);
+
+  /// Undo the above change.
+  void removeDisjunctionChoice(ConstraintLocator *locator) {
+    bool erased = DisjunctionChoices.erase(locator);
+    ASSERT(erased);
   }
 
+
   /// Filter the set of disjunction terms, keeping only those where the
   /// predicate returns \c true.
   ///

lib/Sema/CSSolver.cpp

@@ -302,7 +302,7 @@ void ConstraintSystem::applySolution(const Solution &solution) {
 
   // Register the solution's disjunction choices.
   for (auto &choice : solution.DisjunctionChoices) {
-    DisjunctionChoices.insert(choice);
+    recordDisjunctionChoice(choice.first, choice.second);
   }
 
   // Remember all of the argument/parameter matching choices we made.
@@ -662,7 +662,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
 
   numTypeVariables = cs.TypeVariables.size();
   numFixes = cs.Fixes.size();
-  numDisjunctionChoices = cs.DisjunctionChoices.size();
   numAppliedDisjunctions = cs.AppliedDisjunctions.size();
   numArgumentMatchingChoices = cs.argumentMatchingChoices.size();
   numOpenedTypes = cs.OpenedTypes.size();
@@ -727,9 +726,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  // Remove any disjunction choices.
-  truncate(cs.DisjunctionChoices, numDisjunctionChoices);
-
   // Remove any applied disjunctions.
   truncate(cs.AppliedDisjunctions, numAppliedDisjunctions);
 

lib/Sema/CSTrail.cpp

@@ -151,6 +151,16 @@ SolverTrail::Change::addedFixedRequirement(GenericTypeParamType *GP,
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedDisjunctionChoice(ConstraintLocator *locator,
+                                               unsigned index) {
+  Change result;
+  result.Kind = ChangeKind::RecordedDisjunctionChoice;
+  result.Locator = locator;
+  result.Options = index;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -203,6 +213,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     cs.removeFixedRequirement(FixedRequirement.GP, Options,
                               FixedRequirement.ReqTy);
     break;
+
+  case ChangeKind::RecordedDisjunctionChoice:
+    cs.removeDisjunctionChoice(Locator);
+    break;
   }
 }
 
@@ -313,6 +327,13 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     FixedRequirement.ReqTy->print(out, PO);
     out << ")\n";
     break;
+
+  case ChangeKind::RecordedDisjunctionChoice:
+    out << "(recorded disjunction choice at ";
+    Locator->dump(&cs.getASTContext().SourceMgr, out);
+    out << " index ";
+    out << Options << ")\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -293,6 +293,23 @@ void ConstraintSystem::removeFix(ConstraintFix *fix) {
   Fixes.pop_back();
 }
 
+void ConstraintSystem::recordDisjunctionChoice(
+    ConstraintLocator *locator,
+    unsigned index) {
+  // We shouldn't ever register disjunction choices multiple times.
+  auto inserted = DisjunctionChoices.insert(
+      std::make_pair(locator, index));
+  if (!inserted.second) {
+    ASSERT(inserted.first->second == index);
+    return;
+  }
+
+  if (isRecordingChanges()) {
+    recordChange(SolverTrail::Change::recordedDisjunctionChoice(
+      locator, index));
+  }
+}
+
 /// Retrieve a dynamic result signature for the given declaration.
 static std::tuple<char, ObjCSelector, CanType>
 getDynamicResultSignature(ValueDecl *decl) {