Merge pull request #36575 from hborla/simd-arithmetic-operator-partition

[ConstraintSystem] Treat arithmetic SIMD operators like other generic operators when partitioning an overload set.

Author: hborla

include/swift/AST/KnownProtocols.def

@@ -69,6 +69,7 @@ PROTOCOL(Error)
 PROTOCOL_(ErrorCodeProtocol)
 PROTOCOL(OptionSet)
 PROTOCOL(CaseIterable)
+PROTOCOL(SIMD)
 PROTOCOL(SIMDScalar)
 PROTOCOL(BinaryInteger)
 PROTOCOL(RangeReplaceableCollection)

lib/IRGen/GenMeta.cpp

@@ -5251,6 +5251,7 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::Hashable:
   case KnownProtocolKind::CaseIterable:
   case KnownProtocolKind::Comparable:
+  case KnownProtocolKind::SIMD:
   case KnownProtocolKind::SIMDScalar:
   case KnownProtocolKind::BinaryInteger:
   case KnownProtocolKind::ObjectiveCBridgeable:

lib/Sema/CSSolver.cpp

@@ -1946,6 +1946,7 @@ void DisjunctionChoiceProducer::partitionGenericOperators(
   SmallVector<unsigned, 4> concreteOverloads;
   SmallVector<unsigned, 4> numericOverloads;
   SmallVector<unsigned, 4> sequenceOverloads;
+  SmallVector<unsigned, 4> simdOverloads;
   SmallVector<unsigned, 4> otherGenericOverloads;
 
   auto refinesOrConformsTo = [&](NominalTypeDecl *nominal, KnownProtocolKind kind) -> bool {
@@ -1967,7 +1968,10 @@ void DisjunctionChoiceProducer::partitionGenericOperators(
     unsigned index = *iter;
     auto *decl = Choices[index]->getOverloadChoice().getDecl();
     auto *nominal = decl->getDeclContext()->getSelfNominalTypeDecl();
-    if (!decl->getInterfaceType()->is<GenericFunctionType>()) {
+
+    if (isSIMDOperator(decl)) {
+      simdOverloads.push_back(index);
+    } else if (!decl->getInterfaceType()->is<GenericFunctionType>()) {
       concreteOverloads.push_back(index);
     } else if (refinesOrConformsTo(nominal, KnownProtocolKind::AdditiveArithmetic)) {
       numericOverloads.push_back(index);
@@ -2024,11 +2028,20 @@ void DisjunctionChoiceProducer::partitionGenericOperators(
       sequenceOverloads.clear();
       break;
     }
+
+    if (TypeChecker::conformsToKnownProtocol(
+            argType, KnownProtocolKind::SIMD,
+            CS.DC->getParentModule())) {
+      first = std::copy(simdOverloads.begin(), simdOverloads.end(), first);
+      simdOverloads.clear();
+      break;
+    }
   }
 
   first = std::copy(otherGenericOverloads.begin(), otherGenericOverloads.end(), first);
   first = std::copy(numericOverloads.begin(), numericOverloads.end(), first);
   first = std::copy(sequenceOverloads.begin(), sequenceOverloads.end(), first);
+  first = std::copy(simdOverloads.begin(), simdOverloads.end(), first);
 }
 
 void DisjunctionChoiceProducer::partitionDisjunction(
@@ -2113,7 +2126,8 @@ void DisjunctionChoiceProducer::partitionDisjunction(
   }
 
   // Partition SIMD operators.
-  if (isOperatorDisjunction(Disjunction)) {
+  if (isOperatorDisjunction(Disjunction) &&
+      !Choices[0]->getOverloadChoice().getName().getBaseIdentifier().isArithmeticOperator()) {
     forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
       if (isSIMDOperator(constraint->getOverloadChoice().getDecl())) {
         simdOperators.push_back(index);