Merge pull request swiftlang#19962 from rudkx/refactor-disjunction-partitioning

[ConstraintSystem] Refactor disjunction partitioning

Author: rudkx

lib/Sema/CSSolver.cpp

@@ -1683,6 +1683,68 @@ getOperatorDesignatedNominalTypes(Constraint *bindOverload) {
   return operatorDecl->getDesignatedNominalTypes();
 }
 
+void ConstraintSystem::partitionForDesignatedTypes(
+    ArrayRef<Constraint *> Choices, ConstraintMatchLoop forEachChoice,
+    PartitionAppendCallback appendPartition) {
+
+  auto designatedNominalTypes = getOperatorDesignatedNominalTypes(Choices[0]);
+  if (designatedNominalTypes.empty())
+    return;
+
+  SmallVector<SmallVector<unsigned, 4>, 4> definedInDesignatedType;
+  SmallVector<SmallVector<unsigned, 4>, 4> definedInExtensionOfDesignatedType;
+
+  auto examineConstraint =
+    [&](unsigned constraintIndex, Constraint *constraint) -> bool {
+    auto *decl = constraint->getOverloadChoice().getDecl();
+    auto *funcDecl = cast<FuncDecl>(decl);
+
+    auto *parentDC = funcDecl->getParent();
+    auto *parentDecl = parentDC->getAsDecl();
+
+    if (parentDC->isExtensionContext())
+      parentDecl = cast<ExtensionDecl>(parentDecl)->getExtendedNominal();
+
+    for (auto designatedTypeIndex : indices(designatedNominalTypes)) {
+      auto *designatedNominal =
+        designatedNominalTypes[designatedTypeIndex];
+
+      if (parentDecl != designatedNominal)
+        continue;
+
+      auto &constraints =
+          parentDC->isExtensionContext()
+              ? definedInExtensionOfDesignatedType[designatedTypeIndex]
+              : definedInDesignatedType[designatedTypeIndex];
+
+      constraints.push_back(constraintIndex);
+      return true;
+    }
+
+    return false;
+  };
+
+  definedInDesignatedType.resize(designatedNominalTypes.size());
+  definedInExtensionOfDesignatedType.resize(designatedNominalTypes.size());
+
+  forEachChoice(Choices, examineConstraint);
+
+  // Now collect the overload choices that are defined within the type
+  // that was designated in the operator declaration.
+  // Add partitions for each of the overloads we found in types that
+  // were designated as part of the operator declaration.
+  for (auto designatedTypeIndex : indices(designatedNominalTypes)) {
+    if (designatedTypeIndex < definedInDesignatedType.size()) {
+      auto &primary = definedInDesignatedType[designatedTypeIndex];
+      appendPartition(primary);
+    }
+    if (designatedTypeIndex < definedInExtensionOfDesignatedType.size()) {
+      auto &secondary = definedInExtensionOfDesignatedType[designatedTypeIndex];
+      appendPartition(secondary);
+    }
+  }
+}
+
 void ConstraintSystem::partitionDisjunction(
     ArrayRef<Constraint *> Choices, SmallVectorImpl<unsigned> &Ordering,
     SmallVectorImpl<unsigned> &PartitionBeginning) {
@@ -1702,20 +1764,14 @@ void ConstraintSystem::partitionDisjunction(
     return;
   }
 
-  SmallVector<unsigned, 4> disabled;
-  SmallVector<unsigned, 4> unavailable;
-  SmallVector<unsigned, 4> globalScope;
-  SmallVector<SmallVector<unsigned, 4>, 4> definedInDesignatedType;
-  SmallVector<SmallVector<unsigned, 4>, 4> definedInExtensionOfDesignatedType;
-  SmallVector<unsigned, 4> everythingElse;
   SmallSet<Constraint *, 16> taken;
 
   // 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.
-  auto forEachChoice =
+  ConstraintMatchLoop forEachChoice =
       [&](ArrayRef<Constraint *>,
-          llvm::function_ref<bool(unsigned index, Constraint *)> fn) {
+          std::function<bool(unsigned index, Constraint *)> fn) {
         for (auto index : indices(Choices)) {
           auto *constraint = Choices[index];
           if (taken.count(constraint))
@@ -1729,6 +1785,13 @@ void ConstraintSystem::partitionDisjunction(
         }
       };
 
+  // First collect some things that we'll generally put near the end
+  // of the partitioning.
+
+  SmallVector<unsigned, 4> disabled;
+  SmallVector<unsigned, 4> unavailable;
+  SmallVector<unsigned, 4> globalScope;
+
   // First collect disabled constraints.
   forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
     if (!constraint->isDisabled())
@@ -1765,77 +1828,27 @@ void ConstraintSystem::partitionDisjunction(
     return true;
   });
 
-  // Now collect the overload choices that are defined within the type
-  // that was designated in the operator declaration.
-  auto designatedNominalTypes = getOperatorDesignatedNominalTypes(Choices[0]);
-  if (!designatedNominalTypes.empty()) {
-    forEachChoice(
-        Choices, [&](unsigned constraintIndex, Constraint *constraint) -> bool {
-          auto *decl = constraint->getOverloadChoice().getDecl();
-          auto *funcDecl = cast<FuncDecl>(decl);
-
-          auto *parentDecl = funcDecl->getParent()->getAsDecl();
-          for (auto designatedTypeIndex : indices(designatedNominalTypes)) {
-            auto *designatedNominal =
-                designatedNominalTypes[designatedTypeIndex];
-            if (parentDecl == designatedNominal) {
-              if (designatedTypeIndex >= definedInDesignatedType.size())
-                definedInDesignatedType.resize(designatedTypeIndex + 1);
-              auto &constraints = definedInDesignatedType[designatedTypeIndex];
-              constraints.push_back(constraintIndex);
-              return true;
-            }
-
-            if (auto *extensionDecl = dyn_cast<ExtensionDecl>(parentDecl)) {
-              parentDecl = extensionDecl->getExtendedNominal();
-              if (parentDecl == designatedNominal) {
-                if (designatedTypeIndex >=
-                    definedInExtensionOfDesignatedType.size())
-                  definedInExtensionOfDesignatedType.resize(
-                      designatedTypeIndex + 1);
-
-                auto &constraints =
-                    definedInExtensionOfDesignatedType[designatedTypeIndex];
-                constraints.push_back(constraintIndex);
-                return true;
-              }
-            }
-          }
+  // Local function to create the next partition based on the options
+  // passed in.
+  PartitionAppendCallback appendPartition =
+      [&](SmallVectorImpl<unsigned> &options) {
+        if (options.size()) {
+          PartitionBeginning.push_back(Ordering.size());
+          Ordering.insert(Ordering.end(), options.begin(), options.end());
+        }
+      };
 
-          return false;
-        });
-  }
+  partitionForDesignatedTypes(Choices, forEachChoice, appendPartition);
 
+  SmallVector<unsigned, 4> everythingElse;
   // 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.
-  auto appendPartition = [&](SmallVectorImpl<unsigned> &options) {
-    if (options.size()) {
-      PartitionBeginning.push_back(Ordering.size());
-      Ordering.insert(Ordering.end(), options.begin(), options.end());
-    }
-  };
-
-  // Now create the partitioning based on what was collected.
-
-  // First we'll add partitions for each of the overloads we found in
-  // types that were designated as part of the operator declaration.
-  for (auto designatedTypeIndex : indices(designatedNominalTypes)) {
-    if (designatedTypeIndex < definedInDesignatedType.size()) {
-      auto &primary = definedInDesignatedType[designatedTypeIndex];
-      appendPartition(primary);
-    }
-    if (designatedTypeIndex < definedInExtensionOfDesignatedType.size()) {
-      auto &secondary = definedInExtensionOfDesignatedType[designatedTypeIndex];
-      appendPartition(secondary);
-    }
-  }
   appendPartition(everythingElse);
+
+  // Now create the remaining partitions from what we previously collected.
   appendPartition(globalScope);
   appendPartition(unavailable);
   appendPartition(disabled);

lib/Sema/ConstraintSystem.h

@@ -3221,6 +3221,19 @@ class ConstraintSystem {
 
   bool haveTypeInformationForAllArguments(FunctionType *fnType);
 
+  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 a disjunction based on those that match
+  // the designated types for the operator that the disjunction was
+  // formed for.
+  void partitionForDesignatedTypes(ArrayRef<Constraint *> Choices,
+                                   ConstraintMatchLoop forEachChoice,
+                                   PartitionAppendCallback appendPartition);
+
   // 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.