Merge pull request swiftlang#36911 from xedin/refactor-common-operator-identification-logic

xedin · web-flow · commit 9e159defe8ec · 2021-04-14T15:43:39.000-07:00
[ConstraintSystem] NFC: Factor operator overload identification and partitioning.
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -4983,27 +4983,6 @@ class ConstraintSystem {
     return getExpressionTooComplex(solutionMemory);
   }
 
-  typedef std::function<bool(unsigned index, Constraint *)> ConstraintMatcher;
-  typedef std::function<void(ArrayRef<Constraint *>, ConstraintMatcher)>
-      ConstraintMatchLoop;
-  typedef std::function<void(SmallVectorImpl<unsigned> &options)>
-      PartitionAppendCallback;
-
-  // 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(ArrayRef<Constraint *> Choices,
-                            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
-  /// 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.
   const FunctionType *getAppliedDisjunctionArgumentFunction(const Constraint *disjunction) {
@@ -5328,6 +5307,8 @@ Type isRawRepresentable(ConstraintSystem &cs, Type type,
 Type getDynamicSelfReplacementType(Type baseObjTy, const ValueDecl *member,
                                    ConstraintLocator *memberLocator);
 
+ValueDecl *getOverloadChoiceDecl(Constraint *choice);
+
 class DisjunctionChoice {
   ConstraintSystem &CS;
   unsigned Index;
@@ -5353,7 +5334,7 @@ class DisjunctionChoice {
   }
 
   bool isUnavailable() const {
-    if (auto *decl = getDecl(Choice))
+    if (auto *decl = getOverloadChoiceDecl(Choice))
       return CS.isDeclUnavailable(decl, Choice->getLocator());
     return false;
   }
@@ -5381,23 +5362,12 @@ class DisjunctionChoice {
   void propagateConversionInfo(ConstraintSystem &cs) const;
 
   static ValueDecl *getOperatorDecl(Constraint *choice) {
-    auto *decl = getDecl(choice);
+    auto *decl = getOverloadChoiceDecl(choice);
     if (!decl)
       return nullptr;
 
     return decl->isOperator() ? decl : nullptr;
   }
-
-  static ValueDecl *getDecl(Constraint *constraint) {
-    if (constraint->getKind() != ConstraintKind::BindOverload)
-      return nullptr;
-
-    auto choice = constraint->getOverloadChoice();
-    if (choice.getKind() != OverloadChoiceKind::Decl)
-      return nullptr;
-
-    return choice.getDecl();
-  }
 };
 
 class TypeVariableBinding {
@@ -5560,7 +5530,7 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
     assert(!disjunction->shouldRememberChoice() || disjunction->getLocator());
 
     // Order and partition the disjunction choices.
-    CS.partitionDisjunction(Choices, Ordering, PartitionBeginning);
+    partitionDisjunction(Ordering, PartitionBeginning);
   }
 
   void setNeedsGenericOperatorOrdering(bool flag) {
@@ -5586,16 +5556,27 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
     if (needsGenericOperatorOrdering && choice.isGenericOperator()) {
       unsigned nextPartitionIndex = (PartitionIndex < PartitionBeginning.size() ?
                                      PartitionBeginning[PartitionIndex] : Ordering.size());
-      CS.partitionGenericOperators(Choices,
-                                   Ordering.begin() + currIndex,
-                                   Ordering.begin() + nextPartitionIndex,
-                                   Disjunction->getLocator());
+      partitionGenericOperators(Ordering.begin() + currIndex,
+                                Ordering.begin() + nextPartitionIndex);
       needsGenericOperatorOrdering = false;
     }
 
     return DisjunctionChoice(CS, currIndex, Choices[Ordering[currIndex]],
                              IsExplicitConversion, isBeginningOfPartition);
   }
+
+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);
+
+  /// 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(SmallVectorImpl<unsigned>::iterator first,
+                                 SmallVectorImpl<unsigned>::iterator last);
 };
 
 /// Determine whether given type is a known one
@@ -5620,6 +5601,10 @@ void performSyntacticDiagnosticsForTarget(
 /// generic requirement and if so return that type or null type otherwise.
 Type getConcreteReplacementForProtocolSelfType(ValueDecl *member);
 
+/// Determine whether given disjunction constraint represents a set
+/// of operator overload choices.
+bool isOperatorDisjunction(Constraint *disjunction);
+
 } // end namespace constraints
 
 template<typename ...Args>
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -1734,18 +1734,6 @@ Constraint *ConstraintSystem::getUnboundBindOverloadDisjunction(
   return result->first;
 }
 
-static bool isOperatorBindOverload(Constraint *bindOverload) {
-  if (bindOverload->getKind() != ConstraintKind::BindOverload)
-    return false;
-
-  auto choice = bindOverload->getOverloadChoice();
-  if (!choice.isDecl())
-    return false;
-
-  auto *funcDecl = dyn_cast<FuncDecl>(choice.getDecl());
-  return funcDecl && funcDecl->getOperatorDecl();
-}
-
 // Performance hack: if there are two generic overloads, and one is
 // more specialized than the other, prefer the more-specialized one.
 static Constraint *tryOptimizeGenericDisjunction(
@@ -1878,15 +1866,14 @@ static void existingOperatorBindingsForDisjunction(ConstraintSystem &CS,
   }
 }
 
-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)
+void DisjunctionChoiceProducer::partitionGenericOperators(
+    SmallVectorImpl<unsigned>::iterator first,
+    SmallVectorImpl<unsigned>::iterator last) {
+  auto *argFnType = CS.getAppliedDisjunctionArgumentFunction(Disjunction);
+  if (!isOperatorDisjunction(Disjunction) || !argFnType)
     return;
 
-  auto operatorName = constraints[0]->getOverloadChoice().getName();
+  auto operatorName = Choices[0]->getOverloadChoice().getName();
   if (!operatorName.getBaseIdentifier().isArithmeticOperator())
     return;
 
@@ -1899,7 +1886,8 @@ void ConstraintSystem::partitionGenericOperators(ArrayRef<Constraint *> constrai
     if (!nominal)
       return false;
 
-    auto *protocol = TypeChecker::getProtocol(getASTContext(), SourceLoc(), kind);
+    auto *protocol =
+        TypeChecker::getProtocol(CS.getASTContext(), SourceLoc(), kind);
 
     if (auto *refined = dyn_cast<ProtocolDecl>(nominal))
       return refined->inheritsFrom(protocol);
@@ -1911,7 +1899,7 @@ void ConstraintSystem::partitionGenericOperators(ArrayRef<Constraint *> constrai
   // 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 *decl = Choices[index]->getOverloadChoice().getDecl();
     auto *nominal = decl->getDeclContext()->getSelfNominalTypeDecl();
     if (!decl->getInterfaceType()->is<GenericFunctionType>()) {
       concreteOverloads.push_back(index);
@@ -1926,8 +1914,10 @@ void ConstraintSystem::partitionGenericOperators(ArrayRef<Constraint *> constrai
 
   auto sortPartition = [&](SmallVectorImpl<unsigned> &partition) {
     llvm::sort(partition, [&](unsigned lhs, unsigned rhs) -> bool {
-      auto *declA = dyn_cast<ValueDecl>(constraints[lhs]->getOverloadChoice().getDecl());
-      auto *declB = dyn_cast<ValueDecl>(constraints[rhs]->getOverloadChoice().getDecl());
+      auto *declA =
+          dyn_cast<ValueDecl>(Choices[lhs]->getOverloadChoice().getDecl());
+      auto *declB =
+          dyn_cast<ValueDecl>(Choices[rhs]->getOverloadChoice().getDecl());
 
       return TypeChecker::isDeclRefinementOf(declA, declB);
     });
@@ -1946,19 +1936,22 @@ void ConstraintSystem::partitionGenericOperators(ArrayRef<Constraint *> constrai
   // overload choices first.
   for (auto arg : argFnType->getParams()) {
     auto argType = arg.getPlainType();
-    argType = getFixedTypeRecursive(argType, /*wantRValue=*/true);
+    argType = CS.getFixedTypeRecursive(argType, /*wantRValue=*/true);
 
     if (argType->isTypeVariableOrMember())
       continue;
 
-    if (conformsToKnownProtocol(DC, argType, KnownProtocolKind::AdditiveArithmetic)) {
-      first = std::copy(numericOverloads.begin(), numericOverloads.end(), first);
+    if (conformsToKnownProtocol(CS.DC, argType,
+                                KnownProtocolKind::AdditiveArithmetic)) {
+      first =
+          std::copy(numericOverloads.begin(), numericOverloads.end(), first);
       numericOverloads.clear();
       break;
     }
 
-    if (conformsToKnownProtocol(DC, argType, KnownProtocolKind::Sequence)) {
-      first = std::copy(sequenceOverloads.begin(), sequenceOverloads.end(), first);
+    if (conformsToKnownProtocol(CS.DC, argType, KnownProtocolKind::Sequence)) {
+      first =
+          std::copy(sequenceOverloads.begin(), sequenceOverloads.end(), first);
       sequenceOverloads.clear();
       break;
     }
@@ -1969,17 +1962,23 @@ void ConstraintSystem::partitionGenericOperators(ArrayRef<Constraint *> constrai
   first = std::copy(sequenceOverloads.begin(), sequenceOverloads.end(), first);
 }
 
-void ConstraintSystem::partitionDisjunction(
-    ArrayRef<Constraint *> Choices, SmallVectorImpl<unsigned> &Ordering,
+void DisjunctionChoiceProducer::partitionDisjunction(
+    SmallVectorImpl<unsigned> &Ordering,
     SmallVectorImpl<unsigned> &PartitionBeginning) {
   // Apply a special-case rule for favoring one generic function over
   // another.
-  if (auto favored = tryOptimizeGenericDisjunction(DC, Choices)) {
-    favorConstraint(favored);
+  if (auto favored = tryOptimizeGenericDisjunction(CS.DC, Choices)) {
+    CS.favorConstraint(favored);
   }
 
   SmallSet<Constraint *, 16> taken;
 
+  using ConstraintMatcher = std::function<bool(unsigned index, Constraint *)>;
+  using ConstraintMatchLoop =
+      std::function<void(ArrayRef<Constraint *>, ConstraintMatcher)>;
+  using PartitionAppendCallback =
+      std::function<void(SmallVectorImpl<unsigned> & options)>;
+
   // Local function used to iterate over the untaken choices from the
   // disjunction and use a higher-order function to determine if they
   // should be part of a partition.
@@ -2006,7 +2005,7 @@ void ConstraintSystem::partitionDisjunction(
 
   // Add existing operator bindings to the main partition first. This often
   // helps the solver find a solution fast.
-  existingOperatorBindingsForDisjunction(*this, Choices, everythingElse);
+  existingOperatorBindingsForDisjunction(CS, Choices, everythingElse);
   for (auto index : everythingElse)
     taken.insert(Choices[index]);
 
@@ -2026,7 +2025,7 @@ void ConstraintSystem::partitionDisjunction(
   });
 
   // Then unavailable constraints if we're skipping them.
-  if (!shouldAttemptFixes()) {
+  if (!CS.shouldAttemptFixes()) {
     forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
       if (constraint->getKind() != ConstraintKind::BindOverload)
         return false;
@@ -2036,7 +2035,7 @@ void ConstraintSystem::partitionDisjunction(
       if (!funcDecl)
         return false;
 
-      if (!isDeclUnavailable(funcDecl, constraint->getLocator()))
+      if (!CS.isDeclUnavailable(funcDecl, constraint->getLocator()))
         return false;
 
       unavailable.push_back(index);
@@ -2045,11 +2044,8 @@ void ConstraintSystem::partitionDisjunction(
   }
 
   // Partition SIMD operators.
-  if (isOperatorBindOverload(Choices[0])) {
+  if (isOperatorDisjunction(Disjunction)) {
     forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
-      if (!isOperatorBindOverload(constraint))
-        return false;
-
       if (isSIMDOperator(constraint->getOverloadChoice().getDecl())) {
         simdOperators.push_back(index);
         return true;
@@ -2103,8 +2099,7 @@ Constraint *ConstraintSystem::selectDisjunction() {
         unsigned firstFavored = first->countFavoredNestedConstraints();
         unsigned secondFavored = second->countFavoredNestedConstraints();
 
-        if (!isOperatorBindOverload(first->getNestedConstraints().front()) ||
-            !isOperatorBindOverload(second->getNestedConstraints().front()))
+        if (!isOperatorDisjunction(first) || !isOperatorDisjunction(second))
           return firstActive < secondActive;
 
         if (firstFavored == secondFavored) {
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -5689,3 +5689,19 @@ TypeVarBindingProducer::getDefaultBinding(Constraint *constraint) const {
              ? binding.withType(OptionalType::get(type))
              : binding;
 }
+
+ValueDecl *constraints::getOverloadChoiceDecl(Constraint *choice) {
+  if (choice->getKind() != ConstraintKind::BindOverload)
+    return nullptr;
+  return choice->getOverloadChoice().getDeclOrNull();
+}
+
+bool constraints::isOperatorDisjunction(Constraint *disjunction) {
+  assert(disjunction->getKind() == ConstraintKind::Disjunction);
+
+  auto choices = disjunction->getNestedConstraints();
+  assert(!choices.empty());
+
+  auto *decl = getOverloadChoiceDecl(choices.front());
+  return decl ? decl->isOperator() : false;
+}
