[Constraint solver] Move tryOptimizeGenericDisjunction() into partitioning

This narrow favoring rule makes more sense as part of disjunction
partitioning, because it is not dependent on the use site at all and
should only kick in when other options fail.

Author: DougGregor

lib/Sema/CSSimplify.cpp

@@ -5036,34 +5036,6 @@ Type ConstraintSystem::simplifyAppliedOverloads(
     break;
   }
 
-  // Collect the active overload choices.
-  SmallVector<OverloadChoice, 4> choices;
-  for (auto constraint : disjunction->getNestedConstraints()) {
-    if (constraint->isDisabled())
-      continue;
-    choices.push_back(constraint->getOverloadChoice());
-  }
-
-  // If we can favor one generic result over another, do so.
-  if (auto favoredChoice = tryOptimizeGenericDisjunction(choices)) {
-    unsigned favoredIndex = favoredChoice - choices.data();
-    for (auto constraint : disjunction->getNestedConstraints()) {
-      if (constraint->isDisabled())
-        continue;
-
-      if (favoredIndex == 0) {
-        if (solverState)
-          solverState->favorConstraint(constraint);
-        else
-          constraint->setFavored();
-
-        break;
-      } else {
-        --favoredIndex;
-      }
-    }
-  }
-
   // If there was a constraint that we couldn't reason about, don't use the
   // results of any common-type computations.
   if (hasUnhandledConstraints)

lib/Sema/CSSolver.cpp

@@ -1907,9 +1907,73 @@ void ConstraintSystem::partitionForDesignatedTypes(
   }
 }
 
+// Performance hack: if there are two generic overloads, and one is
+// more specialized than the other, prefer the more-specialized one.
+static Constraint *tryOptimizeGenericDisjunction(
+                                          TypeChecker &tc,
+                                          DeclContext *dc,
+                                          ArrayRef<Constraint *> constraints) {
+  if (constraints.size() != 2 ||
+      constraints[0]->getKind() != ConstraintKind::BindOverload ||
+      constraints[1]->getKind() != ConstraintKind::BindOverload ||
+      constraints[0]->isFavored() ||
+      constraints[1]->isFavored())
+    return nullptr;
+
+  OverloadChoice choiceA = constraints[0]->getOverloadChoice();
+  OverloadChoice choiceB = constraints[1]->getOverloadChoice();
+
+  if (!choiceA.isDecl() || !choiceB.isDecl())
+    return nullptr;
+
+  auto isViable = [](ValueDecl *decl) -> bool {
+    assert(decl);
+
+    auto *AFD = dyn_cast<AbstractFunctionDecl>(decl);
+    if (!AFD || !AFD->isGeneric())
+      return false;
+
+    auto funcType = AFD->getInterfaceType();
+    auto hasAnyOrOptional = funcType.findIf([](Type type) -> bool {
+      if (type->getOptionalObjectType())
+        return true;
+
+      return type->isAny();
+    });
+
+    // If function declaration references `Any` or `Any?` type
+    // let's not attempt it, because it's unclear
+    // without solving which overload is going to be better.
+    return !hasAnyOrOptional;
+  };
+
+  auto *declA = choiceA.getDecl();
+  auto *declB = choiceB.getDecl();
+
+  if (!isViable(declA) || !isViable(declB))
+    return nullptr;
+
+  switch (tc.compareDeclarations(dc, declA, declB)) {
+  case Comparison::Better:
+    return constraints[0];
+
+  case Comparison::Worse:
+    return constraints[1];
+
+  case Comparison::Unordered:
+    return nullptr;
+  }
+}
+
 void ConstraintSystem::partitionDisjunction(
     ArrayRef<Constraint *> Choices, SmallVectorImpl<unsigned> &Ordering,
     SmallVectorImpl<unsigned> &PartitionBeginning) {
+  // Apply a special-case rule for favoring one generic function over
+  // another.
+  if (auto favored = tryOptimizeGenericDisjunction(TC, DC, Choices)) {
+    favorConstraint(favored);
+  }
+
   SmallSet<Constraint *, 16> taken;
 
   // Local function used to iterate over the untaken choices from the

lib/Sema/ConstraintSystem.cpp

@@ -1371,58 +1371,6 @@ ConstraintSystem::getTypeOfMemberReference(
   return { openedType, type };
 }
 
-// Performance hack: if there are two generic overloads, and one is
-// more specialized than the other, prefer the more-specialized one.
-OverloadChoice *ConstraintSystem::tryOptimizeGenericDisjunction(
-                                          ArrayRef<OverloadChoice> choices) {
-  if (choices.size() != 2)
-    return nullptr;
-
-  const auto &choiceA = choices[0];
-  const auto &choiceB = choices[1];
-
-  if (!choiceA.isDecl() || !choiceB.isDecl())
-    return nullptr;
-
-  auto isViable = [](ValueDecl *decl) -> bool {
-    assert(decl);
-
-    auto *AFD = dyn_cast<AbstractFunctionDecl>(decl);
-    if (!AFD || !AFD->isGeneric())
-      return false;
-
-    auto funcType = AFD->getInterfaceType();
-    auto hasAnyOrOptional = funcType.findIf([](Type type) -> bool {
-      if (type->getOptionalObjectType())
-        return true;
-
-      return type->isAny();
-    });
-
-    // If function declaration references `Any` or `Any?` type
-    // let's not attempt it, because it's unclear
-    // without solving which overload is going to be better.
-    return !hasAnyOrOptional;
-  };
-
-  auto *declA = choiceA.getDecl();
-  auto *declB = choiceB.getDecl();
-
-  if (!isViable(declA) || !isViable(declB))
-    return nullptr;
-
-  switch (TC.compareDeclarations(DC, declA, declB)) {
-  case Comparison::Better:
-    return const_cast<OverloadChoice *>(&choiceA);
-
-  case Comparison::Worse:
-    return const_cast<OverloadChoice *>(&choiceB);
-
-  case Comparison::Unordered:
-    return nullptr;
-  }
-}
-
 /// If there are any SIMD operators in the overload set, partition the set so
 /// that the SIMD operators come at the end.
 static ArrayRef<OverloadChoice> partitionSIMDOperators(
@@ -1576,9 +1524,6 @@ void ConstraintSystem::addOverloadSet(Type boundType,
     return;
   }
 
-  if (!favoredChoice)
-    favoredChoice = tryOptimizeGenericDisjunction(choices);
-
   SmallVector<OverloadChoice, 4> scratchChoices;
   choices = partitionSIMDOperators(choices, scratchChoices);
 

lib/Sema/ConstraintSystem.h

@@ -1352,10 +1352,10 @@ class ConstraintSystem {
     /// Favor the given constraint; this change will be rolled back
     /// when we exit the current solver scope.
     void favorConstraint(Constraint *constraint) {
-      if (!constraint->isFavored()) {
-        constraint->setFavored();
-        favoredConstraints.push_back(constraint);
-      }
+      assert(!constraint->isFavored());
+
+      constraint->setFavored();
+      favoredConstraints.push_back(constraint);
     }
 
   private:
@@ -2085,6 +2085,20 @@ class ConstraintSystem {
     removeInactiveConstraint(constraint);
   }
 
+  /// Note that this constraint is "favored" within its disjunction, and
+  /// should be tried first to the exclusion of non-favored constraints in
+  /// the same disjunction.
+  void favorConstraint(Constraint *constraint) {
+    if (constraint->isFavored())
+      return;
+
+    if (solverState) {
+      solverState->favorConstraint(constraint);
+    } else {
+      constraint->setFavored();
+    }
+  }
+
   /// Retrieve the list of inactive constraints.
   ConstraintList &getConstraints() { return InactiveConstraints; }
 
@@ -3187,11 +3201,6 @@ class ConstraintSystem {
 
   Constraint *selectApplyDisjunction();
 
-  /// Look at the set of overload choices to determine if there is a best
-  /// generic overload to favor.
-  OverloadChoice *tryOptimizeGenericDisjunction(
-                                            ArrayRef<OverloadChoice> choices);
-
   /// Solve the system of constraints generated from provided expression.
   ///
   /// \param expr The expression to generate constraints from.