[ConstraintSystem] Cache applied disjunction constraints in the

constraint system to use later in DisjunctionStep.

Author: hborla

include/swift/Sema/ConstraintSystem.h

@@ -2177,6 +2177,11 @@ class ConstraintSystem {
   std::vector<std::pair<ConstraintLocator*, unsigned>>
       DisjunctionChoices;
 
+  /// A map from applied disjunction constraints to the corresponding
+  /// argument function type.
+  llvm::SmallMapVector<ConstraintLocator *, const FunctionType *, 4>
+      AppliedDisjunctions;
+
   /// For locators associated with call expressions, the trailing closure
   /// matching rule that was applied.
   std::vector<std::pair<ConstraintLocator*, TrailingClosureMatching>>
@@ -2669,6 +2674,9 @@ class ConstraintSystem {
     /// The length of \c DisjunctionChoices.
     unsigned numDisjunctionChoices;
 
+    /// The length of \c AppliedDisjunctions.
+    unsigned numAppliedDisjunctions;
+
     /// The length of \c trailingClosureMatchingChoices;
     unsigned numTrailingClosureMatchingChoices;
 
@@ -3923,7 +3931,6 @@ class ConstraintSystem {
       llvm::function_ref<void(unsigned int, Type, ConstraintLocator *)>
           verifyThatArgumentIsHashable);
 
-public:
   /// Describes a direction of optional wrapping, either increasing optionality
   /// or decreasing optionality.
   enum class OptionalWrappingDirection {
@@ -3951,7 +3958,6 @@ class ConstraintSystem {
       TypeVariableType *typeVar, OptionalWrappingDirection optionalDirection,
       llvm::function_ref<bool(Constraint *, TypeVariableType *)> predicate);
 
-private:
   /// Attempt to simplify the set of overloads corresponding to a given
   /// function application constraint.
   ///
@@ -5460,6 +5466,13 @@ class ConstraintSystem {
                             SmallVectorImpl<unsigned> &Ordering,
                             SmallVectorImpl<unsigned> &PartitionBeginning);
 
+  // If the given constraint is an applied disjunction, get the argument function
+  // that the disjunction is applied to.
+  const FunctionType *getAppliedDisjunctionArgumentFunction(Constraint *disjunction) {
+    assert(disjunction->getKind() == ConstraintKind::Disjunction);
+    return AppliedDisjunctions[disjunction->getLocator()];
+  }
+
   /// The overload sets that have already been resolved along the current path.
   const llvm::MapVector<ConstraintLocator *, SelectedOverload> &
   getResolvedOverloads() const {

lib/Sema/CSSimplify.cpp

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

lib/Sema/CSSolver.cpp

@@ -471,6 +471,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numFixes = cs.Fixes.size();
   numFixedRequirements = cs.FixedRequirements.size();
   numDisjunctionChoices = cs.DisjunctionChoices.size();
+  numAppliedDisjunctions = cs.AppliedDisjunctions.size();
   numTrailingClosureMatchingChoices = cs.trailingClosureMatchingChoices.size();
   numOpenedTypes = cs.OpenedTypes.size();
   numOpenedExistentialTypes = cs.OpenedExistentialTypes.size();
@@ -526,6 +527,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any disjunction choices.
   truncate(cs.DisjunctionChoices, numDisjunctionChoices);
 
+  // Remove any applied disjunctions.
+  truncate(cs.AppliedDisjunctions, numAppliedDisjunctions);
+
   // Remove any trailing closure matching choices;
   truncate(
       cs.trailingClosureMatchingChoices, numTrailingClosureMatchingChoices);

lib/Sema/CSStep.cpp

@@ -632,6 +632,7 @@ bool DisjunctionStep::shouldSkip(const DisjunctionChoice &choice) const {
   // introduce any conversions (i.e., the score is not worse than the
   // current score), we can skip any generic operators with conformance
   // requirements that are not satisfied by any known argument types.
+  auto argFnType = CS.getAppliedDisjunctionArgumentFunction(Disjunction);
   auto bestScore = getBestScore(Solutions);
   auto bestChoiceNeedsConversions = bestScore && (bestScore > getCurrentScore());
   if (bestScore && !bestChoiceNeedsConversions && choice.isGenericOperator() && argFnType) {

lib/Sema/CSStep.h

@@ -639,8 +639,6 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
   Optional<Score> BestNonGenericScore;
   Optional<std::pair<Constraint *, Score>> LastSolvedChoice;
 
-  FunctionType *argFnType = nullptr;
-
 public:
   DisjunctionStep(ConstraintSystem &cs, Constraint *disjunction,
                   SmallVectorImpl<Solution> &solutions)
@@ -649,29 +647,6 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
     assert(Disjunction->getKind() == ConstraintKind::Disjunction);
     pruneOverloadSet(Disjunction);
     ++cs.solverState->NumDisjunctions;
-
-    // FIXME: This is duplicate (and expensive) work from simplifyAppliedOverloads
-    auto choices = disjunction->getNestedConstraints();
-    auto *typeVar = choices.front()->getFirstType()->getAs<TypeVariableType>();
-    if (!typeVar)
-      return;
-
-    auto result = cs.findConstraintThroughOptionals(
-        typeVar, ConstraintSystem::OptionalWrappingDirection::Unwrap,
-        [&](Constraint *match, TypeVariableType *currentRep) {
-          // Check to see if we have an applicable fn with a type var RHS that
-          // matches the disjunction.
-          if (match->getKind() != ConstraintKind::ApplicableFunction)
-            return false;
-
-          auto *rhsTyVar = match->getSecondType()->getAs<TypeVariableType>();
-          return rhsTyVar && currentRep == cs.getRepresentative(rhsTyVar);
-        });
-
-    if (result) {
-      auto *applicableFn = result->first;
-      argFnType = applicableFn->getFirstType()->castTo<FunctionType>();
-    }
   }
 
   ~DisjunctionStep() override {