[CSStep] Don't favor choices until the disjunction is picked

include/swift/Sema/ConstraintSystem.h

@@ -5213,8 +5213,9 @@ private:
 
   /// Pick a disjunction from the InactiveConstraints list.
   ///
-  /// \returns The selected disjunction.
-  Constraint *selectDisjunction();
+  /// \returns The selected disjunction and a set of it's favored choices.
+  Optional<std::pair<Constraint *, llvm::TinyPtrVector<Constraint *>>>
+  selectDisjunction();
 
   /// Pick a conjunction from the InactiveConstraints list.
   ///
@@ -6143,7 +6144,8 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
 public:
   using Element = DisjunctionChoice;
 
-  DisjunctionChoiceProducer(ConstraintSystem &cs, Constraint *disjunction)
+  DisjunctionChoiceProducer(ConstraintSystem &cs, Constraint *disjunction,
+                            llvm::TinyPtrVector<Constraint *> &favorites)
       : BindingProducer(cs, disjunction->shouldRememberChoice()
                                 ? disjunction->getLocator()
                                 : nullptr),
@@ -6153,6 +6155,11 @@ public:
     assert(disjunction->getKind() == ConstraintKind::Disjunction);
     assert(!disjunction->shouldRememberChoice() || disjunction->getLocator());
 
+    // Mark constraints as favored. This information
+    // is going to be used by partitioner.
+    for (auto *choice : favorites)
+      cs.favorConstraint(choice);
+
     // Order and partition the disjunction choices.
     partitionDisjunction(Ordering, PartitionBeginning);
   }
@@ -6197,8 +6204,9 @@ private:
   // Partition the choices in the disjunction into groups that we will
   // iterate over in an order appropriate to attempt to stop before we
   // have to visit all of the options.
-  void partitionDisjunction(SmallVectorImpl<unsigned> &Ordering,
-                            SmallVectorImpl<unsigned> &PartitionBeginning);
+  void
+  partitionDisjunction(SmallVectorImpl<unsigned> &Ordering,
+                       SmallVectorImpl<unsigned> &PartitionBeginning);
 
   /// Partition the choices in the range \c first to \c last into groups and
   /// order the groups in the best order to attempt based on the argument

lib/Sema/CSOptimizer.cpp

@@ -468,15 +468,16 @@ selectBestBindingDisjunction(ConstraintSystem &cs,
   return firstBindDisjunction;
 }
 
-Constraint *ConstraintSystem::selectDisjunction() {
+Optional<std::pair<Constraint *, llvm::TinyPtrVector<Constraint *>>>
+ConstraintSystem::selectDisjunction() {
   SmallVector<Constraint *, 4> disjunctions;
 
   collectDisjunctions(disjunctions);
   if (disjunctions.empty())
-    return nullptr;
+    return None;
 
   if (auto *disjunction = selectBestBindingDisjunction(*this, disjunctions))
-    return disjunction;
+    return std::make_pair(disjunction, llvm::TinyPtrVector<Constraint *>());
 
   llvm::DenseMap<Constraint *, llvm::TinyPtrVector<Constraint *>> favorings;
   determineBestChoicesInContext(*this, disjunctions, favorings);
@@ -508,14 +509,8 @@ Constraint *ConstraintSystem::selectDisjunction() {
         return firstFavored < secondFavored;
       });
 
-  if (bestDisjunction != disjunctions.end()) {
-    // If selected disjunction has any choices that should be favored
-    // let's record them now.
-    for (auto *choice : favorings[*bestDisjunction])
-      favorConstraint(choice);
+  if (bestDisjunction != disjunctions.end())
+    return std::make_pair(*bestDisjunction, favorings[*bestDisjunction]);
 
-    return *bestDisjunction;
-  }
-
-  return nullptr;
+  return None;
 }

lib/Sema/CSStep.cpp

@@ -360,7 +360,7 @@ StepResult ComponentStep::take(bool prevFailed) {
     }
   });
 
-  auto *disjunction = CS.selectDisjunction();
+  auto disjunction = CS.selectDisjunction();
   auto *conjunction = CS.selectConjunction();
 
   if (CS.isDebugMode()) {
@@ -403,7 +403,8 @@ StepResult ComponentStep::take(bool prevFailed) {
     // Bindings usually happen first, but sometimes we want to prioritize a
     // disjunction or conjunction.
     if (bestBindings) {
-      if (disjunction && !bestBindings->favoredOverDisjunction(disjunction))
+      if (disjunction &&
+          !bestBindings->favoredOverDisjunction(disjunction->first))
         return StepKind::Disjunction;
 
       if (conjunction && !bestBindings->favoredOverConjunction(conjunction))
@@ -426,9 +427,9 @@ StepResult ComponentStep::take(bool prevFailed) {
       return suspend(
           std::make_unique<TypeVariableStep>(*bestBindings, Solutions));
     case StepKind::Disjunction: {
-      CS.retireConstraint(disjunction);
+      CS.retireConstraint(disjunction->first);
       return suspend(
-          std::make_unique<DisjunctionStep>(CS, disjunction, Solutions));
+          std::make_unique<DisjunctionStep>(CS, *disjunction, Solutions));
     }
     case StepKind::Conjunction: {
       CS.retireConstraint(conjunction);

lib/Sema/CSStep.h

@@ -677,9 +677,17 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
   std::optional<std::pair<Constraint *, Score>> LastSolvedChoice;
 
 public:
+  DisjunctionStep(
+      ConstraintSystem &cs,
+      std::pair<Constraint *, llvm::TinyPtrVector<Constraint *>> &disjunction,
+      SmallVectorImpl<Solution> &solutions)
+      : DisjunctionStep(cs, disjunction.first, disjunction.second, solutions) {}
+
   DisjunctionStep(ConstraintSystem &cs, Constraint *disjunction,
+                  llvm::TinyPtrVector<Constraint *> &favoredChoices,
                   SmallVectorImpl<Solution> &solutions)
-      : BindingStep(cs, {cs, disjunction}, solutions), Disjunction(disjunction) {
+      : BindingStep(cs, {cs, disjunction, favoredChoices}, solutions),
+        Disjunction(disjunction) {
     assert(Disjunction->getKind() == ConstraintKind::Disjunction);
     pruneOverloadSet(Disjunction);
     ++cs.solverState->NumDisjunctions;