Sema: Record applied disjunctions in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -61,6 +61,8 @@ class SolverTrail {
     AddedFixedRequirement,
     /// Recorded a disjunction choice.
     RecordedDisjunctionChoice,
+    /// Recorded an applied disjunction.
+    RecordedAppliedDisjunction,
   };
 
   /// A change made to the constraint system.
@@ -177,6 +179,9 @@ class SolverTrail {
     static Change recordedDisjunctionChoice(ConstraintLocator *locator,
                                             unsigned index);
 
+    /// Create a change that recorded an applied disjunction.
+    static Change recordedAppliedDisjunction(ConstraintLocator *locator);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -1526,6 +1526,10 @@ class Solution {
   /// which informs constraint application.
   llvm::DenseMap<ConstraintLocator *, unsigned> DisjunctionChoices;
 
+  /// A map from applied disjunction constraints to the corresponding
+  /// argument function type.
+  llvm::DenseMap<ConstraintLocator *, FunctionType *> AppliedDisjunctions;
+
   /// The set of opened types for a given locator.
   llvm::DenseMap<ConstraintLocator *, ArrayRef<OpenedType>> OpenedTypes;
 
@@ -2332,7 +2336,7 @@ class ConstraintSystem {
 
   /// A map from applied disjunction constraints to the corresponding
   /// argument function type.
-  llvm::SmallMapVector<ConstraintLocator *, const FunctionType *, 4>
+  llvm::SmallDenseMap<ConstraintLocator *, FunctionType *, 4>
       AppliedDisjunctions;
 
   /// For locators associated with call expressions, the trailing closure
@@ -2879,9 +2883,6 @@ class ConstraintSystem {
     /// FIXME: Remove this.
     unsigned numFixes;
 
-    /// The length of \c AppliedDisjunctions.
-    unsigned numAppliedDisjunctions;
-
     /// The length of \c argumentMatchingChoices.
     unsigned numArgumentMatchingChoices;
 
@@ -5325,6 +5326,15 @@ class ConstraintSystem {
     ASSERT(erased);
   }
 
+  /// Record applied disjunction and add a change to the trail.
+  void recordAppliedDisjunction(ConstraintLocator *locator,
+                                FunctionType *type);
+
+  /// Undo the above change.
+  void removeAppliedDisjunction(ConstraintLocator *locator) {
+    bool erased = AppliedDisjunctions.erase(locator);
+    ASSERT(erased);
+  }
 
   /// Filter the set of disjunction terms, keeping only those where the
   /// predicate returns \c true.
@@ -5629,9 +5639,12 @@ class ConstraintSystem {
 
   // If the given constraint is an applied disjunction, get the argument function
   // that the disjunction is applied to.
-  const FunctionType *getAppliedDisjunctionArgumentFunction(const Constraint *disjunction) {
+  FunctionType *getAppliedDisjunctionArgumentFunction(const Constraint *disjunction) {
     assert(disjunction->getKind() == ConstraintKind::Disjunction);
-    return AppliedDisjunctions[disjunction->getLocator()];
+    auto found = AppliedDisjunctions.find(disjunction->getLocator());
+    if (found == AppliedDisjunctions.end())
+      return nullptr;
+    return found->second;
   }
 
   /// The overload sets that have already been resolved along the current path.

lib/Sema/CSSimplify.cpp

@@ -12823,7 +12823,7 @@ bool ConstraintSystem::simplifyAppliedOverloads(
   auto *applicableFn = result->first;
   auto *fnTypeVar = applicableFn->getSecondType()->castTo<TypeVariableType>();
   auto argFnType = applicableFn->getFirstType()->castTo<FunctionType>();
-  AppliedDisjunctions[disjunction->getLocator()] = argFnType;
+  recordAppliedDisjunction(disjunction->getLocator(), argFnType);
   return simplifyAppliedOverloadsImpl(disjunction, fnTypeVar, argFnType,
                                       /*numOptionalUnwraps*/ result->second,
                                       applicableFn->getLocator());
@@ -12843,7 +12843,7 @@ bool ConstraintSystem::simplifyAppliedOverloads(
   if (!disjunction)
     return false;
 
-  AppliedDisjunctions[disjunction->getLocator()] = argFnType;
+  recordAppliedDisjunction(disjunction->getLocator(), argFnType);
   return simplifyAppliedOverloadsImpl(disjunction, fnTypeVar, argFnType,
                                       numOptionalUnwraps, locator);
 }

lib/Sema/CSSolver.cpp

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

lib/Sema/CSTrail.cpp

@@ -161,6 +161,14 @@ SolverTrail::Change::recordedDisjunctionChoice(ConstraintLocator *locator,
   return result;
 }
 
+SolverTrail::Change
+SolverTrail::Change::recordedAppliedDisjunction(ConstraintLocator *locator) {
+  Change result;
+  result.Kind = ChangeKind::RecordedAppliedDisjunction;
+  result.Locator = locator;
+  return result;
+}
+
 void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   auto &cg = cs.getConstraintGraph();
 
@@ -217,6 +225,10 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
   case ChangeKind::RecordedDisjunctionChoice:
     cs.removeDisjunctionChoice(Locator);
     break;
+
+  case ChangeKind::RecordedAppliedDisjunction:
+    cs.removeAppliedDisjunction(Locator);
+    break;
   }
 }
 
@@ -334,6 +346,12 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     out << " index ";
     out << Options << ")\n";
     break;
+
+  case ChangeKind::RecordedAppliedDisjunction:
+    out << "(recorded applied disjunction at ";
+    Locator->dump(&cs.getASTContext().SourceMgr, out);
+    out << ")\n";
+    break;
   }
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -310,6 +310,18 @@ void ConstraintSystem::recordDisjunctionChoice(
   }
 }
 
+void ConstraintSystem::recordAppliedDisjunction(
+    ConstraintLocator *locator, FunctionType *fnType) {
+  // We shouldn't ever register disjunction choices multiple times.
+  auto inserted = AppliedDisjunctions.insert(
+      std::make_pair(locator, fnType));
+  if (inserted.second) {
+    if (isRecordingChanges()) {
+      recordChange(SolverTrail::Change::recordedAppliedDisjunction(locator));
+    }
+  }
+}
+
 /// Retrieve a dynamic result signature for the given declaration.
 static std::tuple<char, ObjCSelector, CanType>
 getDynamicResultSignature(ValueDecl *decl) {