[ConstraintSystem] Only do the work of partitioning the generic operator

hborla · hborla · commit a3957b987ea1 · 2021-01-14T13:49:57.000-08:00
overloads if generic operators are not going to be skipped.
diff --git a/include/swift/AST/KnownProtocols.def b/include/swift/AST/KnownProtocols.def
@@ -69,7 +69,6 @@ PROTOCOL(Error)
 PROTOCOL_(ErrorCodeProtocol)
 PROTOCOL(OptionSet)
 PROTOCOL(CaseIterable)
-PROTOCOL(SIMD)
 PROTOCOL(SIMDScalar)
 PROTOCOL(BinaryInteger)
 
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -5464,8 +5464,15 @@ class ConstraintSystem {
   // have to visit all of the options.
   void partitionDisjunction(ArrayRef<Constraint *> Choices,
                             SmallVectorImpl<unsigned> &Ordering,
-                            SmallVectorImpl<unsigned> &PartitionBeginning,
-                            ConstraintLocator *locator);
+                            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
+  /// function type that the operator is applied to.
+  void partitionGenericOperators(ArrayRef<Constraint *> Choices,
+                                 SmallVectorImpl<unsigned>::iterator first,
+                                 SmallVectorImpl<unsigned>::iterator last,
+                                 ConstraintLocator *locator);
 
   // If the given constraint is an applied disjunction, get the argument function
   // that the disjunction is applied to.
@@ -5975,8 +5982,12 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
 
   bool IsExplicitConversion;
 
+  Constraint *Disjunction;
+
   unsigned Index = 0;
 
+  bool needsGenericOperatorOrdering = true;
+
 public:
   using Element = DisjunctionChoice;
 
@@ -5985,22 +5996,17 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
                                 ? disjunction->getLocator()
                                 : nullptr),
         Choices(disjunction->getNestedConstraints()),
-        IsExplicitConversion(disjunction->isExplicitConversion()) {
+        IsExplicitConversion(disjunction->isExplicitConversion()),
+        Disjunction(disjunction) {
     assert(disjunction->getKind() == ConstraintKind::Disjunction);
     assert(!disjunction->shouldRememberChoice() || disjunction->getLocator());
 
     // Order and partition the disjunction choices.
-    CS.partitionDisjunction(Choices, Ordering, PartitionBeginning, disjunction->getLocator());
+    CS.partitionDisjunction(Choices, Ordering, PartitionBeginning);
   }
 
-  DisjunctionChoiceProducer(ConstraintSystem &cs,
-                            ArrayRef<Constraint *> choices,
-                            ConstraintLocator *locator, bool explicitConversion)
-      : BindingProducer(cs, locator), Choices(choices),
-        IsExplicitConversion(explicitConversion) {
-
-    // Order and partition the disjunction choices.
-    CS.partitionDisjunction(Choices, Ordering, PartitionBeginning, locator);
+  void setNeedsGenericOperatorOrdering(bool flag) {
+    needsGenericOperatorOrdering = flag;
   }
 
   Optional<Element> operator()() override {
@@ -6015,6 +6021,20 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
 
     ++Index;
 
+    auto choice = DisjunctionChoice(CS, currIndex, Choices[Ordering[currIndex]],
+                                    IsExplicitConversion, isBeginningOfPartition);
+    // Partition the generic operators before producing the first generic
+    // operator disjunction choice.
+    if (needsGenericOperatorOrdering && choice.isGenericOperator()) {
+      unsigned nextPartitionIndex = (PartitionIndex < PartitionBeginning.size() ?
+                                     PartitionBeginning[PartitionIndex] : Ordering.size());
+      CS.partitionGenericOperators(Choices,
+                                   Ordering.begin() + currIndex,
+                                   Ordering.begin() + nextPartitionIndex,
+                                   Disjunction->getLocator());
+      needsGenericOperatorOrdering = false;
+    }
+
     return DisjunctionChoice(CS, currIndex, Choices[Ordering[currIndex]],
                              IsExplicitConversion, isBeginningOfPartition);
   }
diff --git a/lib/IRGen/GenMeta.cpp b/lib/IRGen/GenMeta.cpp
@@ -5123,7 +5123,6 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::Differentiable:
   case KnownProtocolKind::FloatingPoint:
   case KnownProtocolKind::Actor:
-  case KnownProtocolKind::SIMD:
     return SpecialProtocol::None;
   }
 
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -2077,27 +2077,28 @@ static void existingOperatorBindingsForDisjunction(ConstraintSystem &CS,
   }
 }
 
-/// Populates the \c found vector with the indices of the given constraints
-/// that are arithmetic operator choices in the best order to attempt based on
-/// the argument function type that the operator is applied to.
-static void takeArithmeticOperators(ConstraintSystem &CS, const FunctionType *argFnType,
-                                    ArrayRef<Constraint *> constraints,
-                                    SmallVectorImpl<unsigned> &found,
-                                    ConstraintSystem::ConstraintMatchLoop forEachChoice) {
+void ConstraintSystem::partitionGenericOperators(ArrayRef<Constraint *> constraints,
+                                                 SmallVectorImpl<unsigned>::iterator first,
+                                                 SmallVectorImpl<unsigned>::iterator last,
+                                                 ConstraintLocator *locator) {
+  auto *argFnType = AppliedDisjunctions[locator];
+  if (!isOperatorBindOverload(constraints[0]) || !argFnType)
+    return;
+
   auto operatorName = constraints[0]->getOverloadChoice().getName();
   if (!operatorName.getBaseIdentifier().isArithmeticOperator())
     return;
 
+  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 {
     if (!nominal)
       return false;
 
-    auto *protocol = TypeChecker::getProtocol(CS.getASTContext(), SourceLoc(), kind);
+    auto *protocol = TypeChecker::getProtocol(getASTContext(), SourceLoc(), kind);
 
     if (auto *refined = dyn_cast<ProtocolDecl>(nominal))
       return refined->inheritsFrom(protocol);
@@ -2106,29 +2107,21 @@ static void takeArithmeticOperators(ConstraintSystem &CS, const FunctionType *ar
                                                  nominal->getDeclContext());
   };
 
-  // Gather concrete, Numeric, Sequence, and SIMD overloads into separate buckets.
-  forEachChoice(constraints, [&](unsigned index, Constraint *constraint) -> bool {
-    auto *decl = constraint->getOverloadChoice().getDecl();
-
-    if (isSIMDOperator(decl)) {
-      simdOverloads.push_back(index);
-      return true;
-    }
-
+  // Gather Numeric and Sequence overloads into separate buckets.
+  for (auto iter = first; iter != last; ++iter) {
+    unsigned index = *iter;
+    auto *decl = constraints[index]->getOverloadChoice().getDecl();
     auto *nominal = decl->getDeclContext()->getSelfNominalTypeDecl();
     if (!decl->getInterfaceType()->is<GenericFunctionType>()) {
-      // Concrete overloads should always be attempted first.
-      found.push_back(index);
+      concreteOverloads.push_back(index);
     } else if (refinesOrConformsTo(nominal, KnownProtocolKind::AdditiveArithmetic)) {
       numericOverloads.push_back(index);
     } else if (refinesOrConformsTo(nominal, KnownProtocolKind::Sequence)) {
       sequenceOverloads.push_back(index);
     } else {
       otherGenericOverloads.push_back(index);
     }
-
-    return true;
-  });
+  }
 
   auto sortPartition = [&](SmallVectorImpl<unsigned> &partition) {
     llvm::sort(partition, [&](unsigned lhs, unsigned rhs) -> bool {
@@ -2144,6 +2137,9 @@ static void takeArithmeticOperators(ConstraintSystem &CS, const FunctionType *ar
   // subsequent choices that the successful choice is a refinement of.
   sortPartition(sequenceOverloads);
 
+  // Attempt concrete overloads first.
+  first = std::copy(concreteOverloads.begin(), concreteOverloads.end(), first);
+
   // Check if any of the known argument types conform to one of the standard
   // arithmetic protocols. If so, the sovler should attempt the corresponding
   // overload choices first.
@@ -2152,34 +2148,27 @@ static void takeArithmeticOperators(ConstraintSystem &CS, const FunctionType *ar
     if (!argType || argType->hasTypeVariable())
       continue;
 
-    if (conformsToKnownProtocol(CS.DC, argType, KnownProtocolKind::AdditiveArithmetic)) {
-      found.append(numericOverloads.begin(), numericOverloads.end());
+    if (conformsToKnownProtocol(DC, argType, KnownProtocolKind::AdditiveArithmetic)) {
+      first = std::copy(numericOverloads.begin(), numericOverloads.end(), first);
       numericOverloads.clear();
       break;
     }
 
-    if (conformsToKnownProtocol(CS.DC, argType, KnownProtocolKind::Sequence)) {
-      found.append(sequenceOverloads.begin(), sequenceOverloads.end());
+    if (conformsToKnownProtocol(DC, argType, KnownProtocolKind::Sequence)) {
+      first = std::copy(sequenceOverloads.begin(), sequenceOverloads.end(), first);
       sequenceOverloads.clear();
       break;
     }
-
-    if (conformsToKnownProtocol(CS.DC, argType, KnownProtocolKind::SIMD)) {
-      found.append(simdOverloads.begin(), simdOverloads.end());
-      simdOverloads.clear();
-      break;
-    }
   }
 
-  found.append(otherGenericOverloads.begin(), otherGenericOverloads.end());
-  found.append(numericOverloads.begin(), numericOverloads.end());
-  found.append(sequenceOverloads.begin(), sequenceOverloads.end());
-  found.append(simdOverloads.begin(), simdOverloads.end());
+  first = std::copy(otherGenericOverloads.begin(), otherGenericOverloads.end(), first);
+  first = std::copy(numericOverloads.begin(), numericOverloads.end(), first);
+  first = std::copy(sequenceOverloads.begin(), sequenceOverloads.end(), first);
 }
 
 void ConstraintSystem::partitionDisjunction(
     ArrayRef<Constraint *> Choices, SmallVectorImpl<unsigned> &Ordering,
-    SmallVectorImpl<unsigned> &PartitionBeginning, ConstraintLocator *locator) {
+    SmallVectorImpl<unsigned> &PartitionBeginning) {
   // Apply a special-case rule for favoring one generic function over
   // another.
   if (auto favored = tryOptimizeGenericDisjunction(DC, Choices)) {
@@ -2252,11 +2241,8 @@ void ConstraintSystem::partitionDisjunction(
     });
   }
 
-  // Gather arithmetic and SIMD operators.
+  // Partition SIMD operators.
   if (isOperatorBindOverload(Choices[0])) {
-    if (auto *argFnType = AppliedDisjunctions[locator])
-      takeArithmeticOperators(*this, argFnType, Choices, everythingElse, forEachChoice);
-
     forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
       if (!isOperatorBindOverload(constraint))
         return false;
diff --git a/lib/Sema/CSStep.cpp b/lib/Sema/CSStep.cpp
@@ -510,8 +510,15 @@ StepResult DisjunctionStep::resume(bool prevFailed) {
     auto score = getBestScore(Solutions);
 
     if (!choice.isGenericOperator() && choice.isSymmetricOperator()) {
-      if (!BestNonGenericScore || score < BestNonGenericScore)
+      if (!BestNonGenericScore || score < BestNonGenericScore) {
         BestNonGenericScore = score;
+        if (shouldSkipGenericOperators()) {
+          // The disjunction choice producer shouldn't do the work
+          // to partition the generic operator choices if generic
+          // operators are going to be skipped.
+          Producer.setNeedsGenericOperatorOrdering(false);
+        }
+      }
     }
 
     AnySolved = true;
@@ -673,16 +680,8 @@ bool DisjunctionStep::shouldSkip(const DisjunctionChoice &choice) const {
   //        already have a solution involving non-generic operators,
   //        but continue looking for a better non-generic operator
   //        solution.
-  if (BestNonGenericScore && choice.isGenericOperator()) {
-    auto &score = BestNonGenericScore->Data;
-    // Let's skip generic overload choices only in case if
-    // non-generic score indicates that there were no forced
-    // unwrappings of optional(s), no unavailable overload
-    // choices present in the solution, no fixes required,
-    // and there are no non-trivial function conversions.
-    if (score[SK_ForceUnchecked] == 0 && score[SK_Unavailable] == 0 &&
-        score[SK_Fix] == 0 && score[SK_FunctionConversion] == 0)
-      return skip("generic");
+  if (shouldSkipGenericOperators() && choice.isGenericOperator()) {
+    return skip("generic");
   }
 
   return false;
diff --git a/lib/Sema/CSStep.h b/lib/Sema/CSStep.h
@@ -683,6 +683,20 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
   bool shortCircuitDisjunctionAt(Constraint *currentChoice,
                                  Constraint *lastSuccessfulChoice) const;
 
+  bool shouldSkipGenericOperators() const {
+    if (!BestNonGenericScore)
+      return false;
+
+    // Let's skip generic overload choices only in case if
+    // non-generic score indicates that there were no forced
+    // unwrappings of optional(s), no unavailable overload
+    // choices present in the solution, no fixes required,
+    // and there are no non-trivial function conversions.
+    auto &score = BestNonGenericScore->Data;
+    return (score[SK_ForceUnchecked] == 0 && score[SK_Unavailable] == 0 &&
+            score[SK_Fix] == 0 && score[SK_FunctionConversion] == 0);
+  }
+
   /// Attempt to apply given disjunction choice to constraint system.
   /// This action is going to establish "active choice" of this disjunction
   /// to point to a given choice.

Original file line number	Diff line number	Diff line change
`@@ -5123,7 +5123,6 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {`
`5123`	`5123`	`case KnownProtocolKind::Differentiable:`
`5124`	`5124`	`case KnownProtocolKind::FloatingPoint:`
`5125`	`5125`	`case KnownProtocolKind::Actor:`
`5126`		`- case KnownProtocolKind::SIMD:`
`5127`	`5126`	`return SpecialProtocol::None;`
`5128`	`5127`	`}`
`5129`	`5128`