Merge pull request swiftlang#35025 from hborla/generic-overload-search-space

[ConstraintSystem] Implement heuristics for pruning the generic operator overload search space

Author: hborla

include/swift/AST/Identifier.h

@@ -110,6 +110,10 @@ class Identifier {
     return isOperatorSlow();
   }
 
+  bool isArithmeticOperator() const {
+    return is("+") || is("-") || is("*") || is("/") || is("%");
+  }
+
   // Returns whether this is a standard comparison operator,
   // such as '==', '>=' or '!=='.
   bool isStandardComparisonOperator() const {

include/swift/AST/TypeCheckRequests.h

@@ -1900,6 +1900,27 @@ class CompareDeclSpecializationRequest
   bool isCached() const { return true; }
 };
 
+/// Checks whether the first function decl is a refinement of the second,
+/// meaning the two functions have the same structure, and the requirements
+/// of the first are refining the requirements of the second.
+class IsDeclRefinementOfRequest
+    : public SimpleRequest<IsDeclRefinementOfRequest,
+                           bool(ValueDecl *, ValueDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  bool evaluate(Evaluator &evaluator, ValueDecl *declA, ValueDecl *declB) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 /// Checks whether this declaration inherits its superclass' designated and
 /// convenience initializers.
 class InheritsSuperclassInitializersRequest

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -41,6 +41,9 @@ SWIFT_REQUEST(TypeChecker, CodeCompletionFileRequest,
 SWIFT_REQUEST(TypeChecker, CompareDeclSpecializationRequest,
               bool (DeclContext *, ValueDecl *, ValueDecl *, bool), Cached,
               NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, IsDeclRefinementOfRequest,
+              bool (ValueDecl *, ValueDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, CustomAttrTypeRequest,
               Type(CustomAttr *, DeclContext *, CustomAttrTypeKind),
               SeparatelyCached, NoLocationInfo)

include/swift/Sema/ConstraintSystem.h

@@ -2164,6 +2164,11 @@ class ConstraintSystem {
   std::vector<std::pair<ConstraintLocator*, unsigned>>
       DisjunctionChoices;
 
+  /// A map from applied disjunction constraints to the corresponding
+  /// argument function type.
+  llvm::SmallMapVector<ConstraintLocator *, const FunctionType *, 4>
+      AppliedDisjunctions;
+
   /// For locators associated with call expressions, the trailing closure
   /// matching rule that was applied.
   std::vector<std::pair<ConstraintLocator*, TrailingClosureMatching>>
@@ -2653,6 +2658,9 @@ class ConstraintSystem {
     /// The length of \c DisjunctionChoices.
     unsigned numDisjunctionChoices;
 
+    /// The length of \c AppliedDisjunctions.
+    unsigned numAppliedDisjunctions;
+
     /// The length of \c trailingClosureMatchingChoices;
     unsigned numTrailingClosureMatchingChoices;
 
@@ -4954,6 +4962,21 @@ class ConstraintSystem {
                             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(Constraint *disjunction) {
+    assert(disjunction->getKind() == ConstraintKind::Disjunction);
+    return AppliedDisjunctions[disjunction->getLocator()];
+  }
+
   /// The overload sets that have already been resolved along the current path.
   const llvm::MapVector<ConstraintLocator *, SelectedOverload> &
   getResolvedOverloads() const {
@@ -5454,8 +5477,12 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
 
   bool IsExplicitConversion;
 
+  Constraint *Disjunction;
+
   unsigned Index = 0;
 
+  bool needsGenericOperatorOrdering = true;
+
 public:
   using Element = DisjunctionChoice;
 
@@ -5464,22 +5491,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);
   }
 
-  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);
+  void setNeedsGenericOperatorOrdering(bool flag) {
+    needsGenericOperatorOrdering = flag;
   }
 
   Optional<Element> operator()() override {
@@ -5494,6 +5516,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);
   }

lib/Sema/CSGen.cpp

@@ -38,12 +38,7 @@ using namespace swift;
 using namespace swift::constraints;
 
 static bool isArithmeticOperatorDecl(ValueDecl *vd) {
-  return vd && 
-  (vd->getBaseName() == "+" ||
-   vd->getBaseName() == "-" ||
-   vd->getBaseName() == "*" ||
-   vd->getBaseName() == "/" ||
-   vd->getBaseName() == "%");
+  return vd && vd->getBaseIdentifier().isArithmeticOperator();
 }
 
 static bool mergeRepresentativeEquivalenceClasses(ConstraintSystem &CS,

lib/Sema/CSSimplify.cpp

@@ -8919,6 +8919,7 @@ bool ConstraintSystem::simplifyAppliedOverloads(
   auto *applicableFn = result->first;
   auto *fnTypeVar = applicableFn->getSecondType()->castTo<TypeVariableType>();
   auto argFnType = applicableFn->getFirstType()->castTo<FunctionType>();
+  AppliedDisjunctions[disjunction->getLocator()] = argFnType;
   return simplifyAppliedOverloadsImpl(disjunction, fnTypeVar, argFnType,
                                       /*numOptionalUnwraps*/ result->second,
                                       locator);
@@ -8938,6 +8939,8 @@ bool ConstraintSystem::simplifyAppliedOverloads(
       getUnboundBindOverloadDisjunction(fnTypeVar, &numOptionalUnwraps);
   if (!disjunction)
     return false;
+
+  AppliedDisjunctions[disjunction->getLocator()] = argFnType;
   return simplifyAppliedOverloadsImpl(disjunction, fnTypeVar, argFnType,
                                       numOptionalUnwraps, locator);
 }

lib/Sema/CSSolver.cpp

@@ -465,6 +465,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numFixes = cs.Fixes.size();
   numFixedRequirements = cs.FixedRequirements.size();
   numDisjunctionChoices = cs.DisjunctionChoices.size();
+  numAppliedDisjunctions = cs.AppliedDisjunctions.size();
   numTrailingClosureMatchingChoices = cs.trailingClosureMatchingChoices.size();
   numOpenedTypes = cs.OpenedTypes.size();
   numOpenedExistentialTypes = cs.OpenedExistentialTypes.size();
@@ -519,6 +520,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any disjunction choices.
   truncate(cs.DisjunctionChoices, numDisjunctionChoices);
 
+  // Remove any applied disjunctions.
+  truncate(cs.AppliedDisjunctions, numAppliedDisjunctions);
+
   // Remove any trailing closure matching choices;
   truncate(
       cs.trailingClosureMatchingChoices, numTrailingClosureMatchingChoices);
@@ -2063,6 +2067,95 @@ 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)
+    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> otherGenericOverloads;
+
+  auto refinesOrConformsTo = [&](NominalTypeDecl *nominal, KnownProtocolKind kind) -> bool {
+    if (!nominal)
+      return false;
+
+    auto *protocol = TypeChecker::getProtocol(getASTContext(), SourceLoc(), kind);
+
+    if (auto *refined = dyn_cast<ProtocolDecl>(nominal))
+      return refined->inheritsFrom(protocol);
+
+    return (bool)TypeChecker::conformsToProtocol(nominal->getDeclaredType(), protocol,
+                                                 nominal->getDeclContext());
+  };
+
+  // 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>()) {
+      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);
+    }
+  }
+
+  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());
+
+      return TypeChecker::isDeclRefinementOf(declA, declB);
+    });
+  };
+
+  // Sort sequence overloads so that refinements are attempted first.
+  // If the solver finds a solution with an overload, it can then skip
+  // 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.
+  for (auto arg : argFnType->getParams()) {
+    auto argType = arg.getPlainType();
+    if (!argType || argType->hasTypeVariable())
+      continue;
+
+    if (conformsToKnownProtocol(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);
+      sequenceOverloads.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);
+}
+
 void ConstraintSystem::partitionDisjunction(
     ArrayRef<Constraint *> Choices, SmallVectorImpl<unsigned> &Ordering,
     SmallVectorImpl<unsigned> &PartitionBeginning) {
@@ -2153,6 +2246,12 @@ void ConstraintSystem::partitionDisjunction(
     });
   }
 
+  // Gather the remaining options.
+  forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
+    everythingElse.push_back(index);
+    return true;
+  });
+
   // Local function to create the next partition based on the options
   // passed in.
   PartitionAppendCallback appendPartition =
@@ -2163,16 +2262,9 @@ void ConstraintSystem::partitionDisjunction(
         }
       };
 
-  // Gather the remaining options.
-  forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
-    everythingElse.push_back(index);
-    return true;
-  });
   appendPartition(favored);
   appendPartition(everythingElse);
   appendPartition(simdOperators);
-
-  // Now create the remaining partitions from what we previously collected.
   appendPartition(unavailable);
   appendPartition(disabled);