[Constraint solver] Move all partitioning into partitionDisjunction

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
@@ -1932,6 +1996,7 @@ void ConstraintSystem::partitionDisjunction(
   // end of the partitioning.
 
   SmallVector<unsigned, 4> favored;
+  SmallVector<unsigned, 4> simdOperators;
   SmallVector<unsigned, 4> disabled;
   SmallVector<unsigned, 4> unavailable;
 
@@ -1971,6 +2036,22 @@ void ConstraintSystem::partitionDisjunction(
     });
   }
 
+  // Partition SIMD operators.
+  if (!TC.getLangOpts().SolverEnableOperatorDesignatedTypes &&
+      isOperatorBindOverload(Choices[0])) {
+    forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
+      if (!isOperatorBindOverload(constraint))
+        return false;
+
+      if (isSIMDOperator(constraint->getOverloadChoice().getDecl())) {
+        simdOperators.push_back(index);
+        return true;
+      }
+
+      return false;
+    });
+  }
+
   // Local function to create the next partition based on the options
   // passed in.
   PartitionAppendCallback appendPartition =
@@ -1994,6 +2075,7 @@ void ConstraintSystem::partitionDisjunction(
   });
   appendPartition(favored);
   appendPartition(everythingElse);
+  appendPartition(simdOperators);
 
   // Now create the remaining partitions from what we previously collected.
   appendPartition(unavailable);

lib/Sema/ConstraintSystem.cpp

@@ -1371,92 +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(
-                                  ArrayRef<OverloadChoice> choices,
-                                  SmallVectorImpl<OverloadChoice> &scratch) {
-  // If the first element isn't an operator, none of them are.
-  if (!choices[0].isDecl() ||
-      !isa<FuncDecl>(choices[0].getDecl()) ||
-      !cast<FuncDecl>(choices[0].getDecl())->isOperator() ||
-      choices[0].getDecl()->getASTContext().LangOpts
-        .SolverEnableOperatorDesignatedTypes)
-    return choices;
-
-  // Check whether we have any SIMD operators.
-  bool foundSIMDOperator = false;
-  for (const auto &choice : choices) {
-    if (isSIMDOperator(choice.getDecl())) {
-      foundSIMDOperator = true;
-      break;
-    }
-  }
-
-  if (!foundSIMDOperator)
-    return choices;
-
-  scratch.assign(choices.begin(), choices.end());
-  std::stable_partition(scratch.begin(), scratch.end(),
-                        [](const OverloadChoice &choice) {
-                          return !isSIMDOperator(choice.getDecl());
-                        });
-
-  return scratch;
-}
-
 Type ConstraintSystem::getEffectiveOverloadType(const OverloadChoice &overload,
                                                 bool allowMembers,
                                                 DeclContext *useDC) {
@@ -1576,12 +1490,6 @@ void ConstraintSystem::addOverloadSet(Type boundType,
     return;
   }
 
-  if (!favoredChoice)
-    favoredChoice = tryOptimizeGenericDisjunction(choices);
-
-  SmallVector<OverloadChoice, 4> scratchChoices;
-  choices = partitionSIMDOperators(choices, scratchChoices);
-
   SmallVector<Constraint *, 4> overloads;
 
   // As we do for other favored constraints, if a favored overload has been

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.