[ConstraintSystem] If the solver has already found a solution with a

disjunction choice that does not introduce conversions, check to see
if known argument types satisfy generic operator conformance requirements
early, and skip the overload choice if any requirements fail.

This helps the solver avoid exploring way too much search space when
the right solution involves a generic operator, but the argument types
are known up front, such as `collection + collection + collection`.

Author: hborla

include/swift/Sema/ConstraintSystem.h

@@ -3923,6 +3923,7 @@ 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 {
@@ -3950,6 +3951,7 @@ 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.
   ///

lib/Sema/CSStep.cpp

@@ -16,6 +16,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "CSStep.h"
+#include "TypeChecker.h"
 #include "swift/AST/Types.h"
 #include "swift/Sema/ConstraintSystem.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -555,6 +556,40 @@ bool DisjunctionStep::shouldSkip(const DisjunctionChoice &choice) const {
   if (ctx.TypeCheckerOpts.DisableConstraintSolverPerformanceHacks)
     return false;
 
+  // If the solver already found a solution with a choice that did not
+  // 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 bestScore = getBestScore(Solutions);
+  auto bestChoiceNeedsConversions = bestScore && (bestScore > getCurrentScore());
+  if (!bestChoiceNeedsConversions && choice.isGenericOperator() && argFnType) {
+    Constraint *constraint = choice;
+    auto *decl = constraint->getOverloadChoice().getDecl();
+    auto *useDC = constraint->getOverloadUseDC();
+    auto choiceType = CS.getEffectiveOverloadType(constraint->getOverloadChoice(),
+                                                  /*allowMembers=*/true, useDC);
+    auto choiceFnType = choiceType->getAs<FunctionType>();
+    auto genericFnType = decl->getInterfaceType()->getAs<GenericFunctionType>();
+    auto signature = genericFnType->getGenericSignature();
+
+    for (auto argParamPair : llvm::zip(argFnType->getParams(),
+                                       choiceFnType->getParams())) {
+      auto argType = std::get<0>(argParamPair).getPlainType();
+      auto paramType = std::get<1>(argParamPair).getPlainType();
+
+      // Only check argument types with no type variables that will be matched
+      // against a plain type parameter.
+      argType = argType->getCanonicalType()->getWithoutSpecifierType();
+      if (argType->hasTypeVariable() || !paramType->isTypeParameter())
+        continue;
+
+      for (auto *protocol : signature->getRequiredProtocols(paramType)) {
+        if (!TypeChecker::conformsToProtocol(argType, protocol, useDC))
+          return skip("unsatisfied");
+      }
+    }
+  }
+
   // Don't attempt to solve for generic operators if we already have
   // a non-generic solution.
 

lib/Sema/CSStep.h

@@ -639,6 +639,8 @@ 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)
@@ -647,6 +649,29 @@ 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 {
@@ -732,7 +757,9 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
 
   // Figure out which of the solutions has the smallest score.
   static Optional<Score> getBestScore(SmallVectorImpl<Solution> &solutions) {
-    assert(!solutions.empty());
+    if (solutions.empty())
+      return None;
+
     Score bestScore = solutions.front().getFixedScore();
     if (solutions.size() == 1)
       return bestScore;