[Constraint solver] Enable favoring of disjunction constraints during solving

Allow constraints to be favored during solving, and unwound after exiting
that particular solver scope. We're not using this yet.

Author: DougGregor

lib/Sema/CSSimplify.cpp

@@ -4903,7 +4903,6 @@ Type ConstraintSystem::simplifyAppliedOverloads(
     break;
   }
 
-
   // 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

@@ -424,6 +424,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numCheckedConformances = cs.CheckedConformances.size();
   numMissingMembers = cs.MissingMembers.size();
   numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
+  numFavoredConstraints = cs.solverState->getNumFavoredConstraints();
 
   PreviousScore = cs.CurrentScore;
 

lib/Sema/ConstraintSystem.cpp

@@ -1373,17 +1373,16 @@ ConstraintSystem::getTypeOfMemberReference(
 
 // Performance hack: if there are two generic overloads, and one is
 // more specialized than the other, prefer the more-specialized one.
-static void tryOptimizeGenericDisjunction(ConstraintSystem &cs,
-                                          ArrayRef<OverloadChoice> choices,
-                                          OverloadChoice *&favoredChoice) {
-  if (favoredChoice || choices.size() != 2)
-    return;
+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;
+    return nullptr;
 
   auto isViable = [](ValueDecl *decl) -> bool {
     assert(decl);
@@ -1410,22 +1409,17 @@ static void tryOptimizeGenericDisjunction(ConstraintSystem &cs,
   auto *declB = choiceB.getDecl();
 
   if (!isViable(declA) || !isViable(declB))
-    return;
-
-  auto &TC = cs.TC;
-  auto *DC = cs.DC;
+    return nullptr;
 
   switch (TC.compareDeclarations(DC, declA, declB)) {
   case Comparison::Better:
-    favoredChoice = const_cast<OverloadChoice *>(&choiceA);
-    break;
+    return const_cast<OverloadChoice *>(&choiceA);
 
   case Comparison::Worse:
-    favoredChoice = const_cast<OverloadChoice *>(&choiceB);
-    break;
+    return const_cast<OverloadChoice *>(&choiceB);
 
   case Comparison::Unordered:
-    break;
+    return nullptr;
   }
 }
 
@@ -1582,7 +1576,8 @@ void ConstraintSystem::addOverloadSet(Type boundType,
     return;
   }
 
-  tryOptimizeGenericDisjunction(*this, choices, favoredChoice);
+  if (!favoredChoice)
+    favoredChoice = tryOptimizeGenericDisjunction(choices);
 
   SmallVector<OverloadChoice, 4> scratchChoices;
   choices = partitionSIMDOperators(choices, scratchChoices);

lib/Sema/ConstraintSystem.h

@@ -1317,6 +1317,15 @@ class ConstraintSystem {
       disabledConstraints.erase(
           disabledConstraints.begin() + scope->numDisabledConstraints,
           disabledConstraints.end());
+
+      for (unsigned constraintIdx :
+             range(scope->numFavoredConstraints, favoredConstraints.size())) {
+        if (favoredConstraints[constraintIdx]->isFavored())
+          favoredConstraints[constraintIdx]->setFavored(false);
+      }
+      favoredConstraints.erase(
+          favoredConstraints.begin() + scope->numFavoredConstraints,
+          favoredConstraints.end());
     }
 
     /// Check whether constraint system is allowed to form solutions
@@ -1336,6 +1345,19 @@ class ConstraintSystem {
       disabledConstraints.push_back(constraint);
     }
 
+    unsigned getNumFavoredConstraints() const {
+      return favoredConstraints.size();
+    }
+
+    /// 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);
+      }
+    }
+
   private:
     /// The list of constraints that have been retired along the
     /// current path, this list is used in LIFO fashion when constraints
@@ -1358,6 +1380,7 @@ class ConstraintSystem {
       std::tuple<SolverScope *, ConstraintList::iterator, unsigned>, 4> scopes;
 
     SmallVector<Constraint *, 4> disabledConstraints;
+    SmallVector<Constraint *, 4> favoredConstraints;
   };
 
   class CacheExprTypes : public ASTWalker {
@@ -1518,6 +1541,8 @@ class ConstraintSystem {
 
     unsigned numDisabledConstraints;
 
+    unsigned numFavoredConstraints;
+
     /// The previous score.
     Score PreviousScore;
 
@@ -3162,6 +3187,11 @@ 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.