[Constraint system] Move SIMD operator partitioning into disjunction partitioning

Continuing along my path to move application-independent partitioning
rules into the common disjunction partitioning code.

Author: DougGregor

lib/Sema/CSSolver.cpp

@@ -1996,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;
 
@@ -2035,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 =
@@ -2058,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,40 +1371,6 @@ ConstraintSystem::getTypeOfMemberReference(
   return { openedType, type };
 }
 
-/// 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) {
@@ -1524,9 +1490,6 @@ void ConstraintSystem::addOverloadSet(Type boundType,
     return;
   }
 
-  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