Merge pull request #18386 from rudkx/simplify-state-for-disjunction-choice

[ConstraintSystem] Split the main loop attempting disjunction choices…

Author: rudkx

lib/Sema/CSSolver.cpp

@@ -1957,64 +1957,12 @@ Constraint *ConstraintSystem::selectDisjunction() {
 }
 
 bool ConstraintSystem::solveForDisjunctionChoices(
-    Constraint *disjunction, SmallVectorImpl<Solution> &solutions,
-    FreeTypeVariableBinding allowFreeTypeVariables) {
-  assert(disjunction->getKind() == ConstraintKind::Disjunction);
-
-  // Remove this disjunction constraint from the list.
-  auto afterDisjunction = InactiveConstraints.erase(disjunction);
-  CG.removeConstraint(disjunction);
-
-  // Check if selected disjunction has a representative
-  // this might happen when there are multiple binary operators
-  // chained together. If so, disable choices which differ
-  // from currently selected representative.
-  auto pruneOverloadSet = [&](Constraint *disjunction) -> bool {
-    auto *choice = disjunction->getNestedConstraints().front();
-    auto *typeVar = choice->getFirstType()->getAs<TypeVariableType>();
-    if (!typeVar)
-      return false;
-
-    auto *repr = typeVar->getImpl().getRepresentative(nullptr);
-    if (!repr || repr == typeVar)
-      return false;
-
-    bool isPruned = false;
-    for (auto resolved = resolvedOverloadSets; resolved;
-         resolved = resolved->Previous) {
-      if (!resolved->BoundType->isEqual(repr))
-        continue;
-
-      auto &representative = resolved->Choice;
-      if (!representative.isDecl())
-        return false;
-
-      // Disable all of the overload choices which are different from
-      // the one which is currently picked for representative.
-      for (auto *constraint : disjunction->getNestedConstraints()) {
-        auto choice = constraint->getOverloadChoice();
-        if (!choice.isDecl())
-          continue;
-
-        if (choice.getDecl() != representative.getDecl()) {
-          constraint->setDisabled();
-          isPruned = true;
-        }
-      }
-      break;
-    }
-
-    return isPruned;
-  };
-
-  bool hasDisabledChoices = pruneOverloadSet(disjunction);
-
-  Optional<std::pair<DisjunctionChoice, Score>> lastSolvedChoice;
+    ArrayRef<Constraint *> constraints, ConstraintLocator *disjunctionLocator,
+    SmallVectorImpl<Solution> &solutions,
+    FreeTypeVariableBinding allowFreeTypeVariables, bool explicitConversion) {
   Optional<Score> bestNonGenericScore;
+  Optional<std::pair<DisjunctionChoice, Score>> lastSolvedChoice;
 
-  ++solverState->NumDisjunctions;
-  bool explicitConversion = isExplicitConversionConstraint(disjunction);
-  auto constraints = disjunction->getNestedConstraints();
   // Try each of the constraints within the disjunction.
   for (auto index : indices(constraints)) {
     auto currentChoice =
@@ -2052,15 +2000,14 @@ bool ConstraintSystem::solveForDisjunctionChoices(
 
     // If the disjunction requested us to, remember which choice we
     // took for it.
-    if (disjunction->shouldRememberChoice()) {
-      auto locator = disjunction->getLocator();
-      assert(locator && "remembered disjunction doesn't have a locator?");
-      DisjunctionChoices.push_back({locator, index});
+
+    if (disjunctionLocator) {
+      DisjunctionChoices.push_back({disjunctionLocator, index});
 
       // Implicit unwraps of optionals are worse solutions than those
       // not involving implicit unwraps.
-      if (!locator->getPath().empty()) {
-        auto kind = locator->getPath().back().getKind();
+      if (!disjunctionLocator->getPath().empty()) {
+        auto kind = disjunctionLocator->getPath().back().getKind();
         if (kind == ConstraintLocator::ImplicitlyUnwrappedDisjunctionChoice ||
             kind == ConstraintLocator::DynamicLookupResult) {
           assert(index == 0 || index == 1);
@@ -2086,6 +2033,71 @@ bool ConstraintSystem::solveForDisjunctionChoices(
     }
   }
 
+  return !bool(lastSolvedChoice.hasValue());
+}
+
+bool ConstraintSystem::solveForDisjunction(
+    Constraint *disjunction, SmallVectorImpl<Solution> &solutions,
+    FreeTypeVariableBinding allowFreeTypeVariables) {
+  assert(disjunction->getKind() == ConstraintKind::Disjunction);
+
+  // Remove this disjunction constraint from the list.
+  auto afterDisjunction = InactiveConstraints.erase(disjunction);
+  CG.removeConstraint(disjunction);
+
+  // Check if selected disjunction has a representative
+  // this might happen when there are multiple binary operators
+  // chained together. If so, disable choices which differ
+  // from currently selected representative.
+  auto pruneOverloadSet = [&](Constraint *disjunction) -> bool {
+    auto *choice = disjunction->getNestedConstraints().front();
+    auto *typeVar = choice->getFirstType()->getAs<TypeVariableType>();
+    if (!typeVar)
+      return false;
+
+    auto *repr = typeVar->getImpl().getRepresentative(nullptr);
+    if (!repr || repr == typeVar)
+      return false;
+
+    bool isPruned = false;
+    for (auto resolved = resolvedOverloadSets; resolved;
+         resolved = resolved->Previous) {
+      if (!resolved->BoundType->isEqual(repr))
+        continue;
+
+      auto &representative = resolved->Choice;
+      if (!representative.isDecl())
+        return false;
+
+      // Disable all of the overload choices which are different from
+      // the one which is currently picked for representative.
+      for (auto *constraint : disjunction->getNestedConstraints()) {
+        auto choice = constraint->getOverloadChoice();
+        if (!choice.isDecl())
+          continue;
+
+        if (choice.getDecl() != representative.getDecl()) {
+          constraint->setDisabled();
+          isPruned = true;
+        }
+      }
+      break;
+    }
+
+    return isPruned;
+  };
+
+  bool hasDisabledChoices = pruneOverloadSet(disjunction);
+
+  ++solverState->NumDisjunctions;
+  auto *locator =
+      disjunction->shouldRememberChoice() ? disjunction->getLocator() : nullptr;
+  assert(!disjunction->shouldRememberChoice() || disjunction->getLocator());
+
+  auto noSolutions = solveForDisjunctionChoices(
+      disjunction->getNestedConstraints(), locator, solutions,
+      allowFreeTypeVariables, isExplicitConversionConstraint(disjunction));
+
   if (hasDisabledChoices) {
     // Re-enable previously disabled overload choices.
     for (auto *choice : disjunction->getNestedConstraints()) {
@@ -2098,7 +2110,7 @@ bool ConstraintSystem::solveForDisjunctionChoices(
   InactiveConstraints.insert(afterDisjunction, disjunction);
   CG.addConstraint(disjunction);
 
-  return bool(lastSolvedChoice.hasValue());
+  return noSolutions;
 }
 
 bool ConstraintSystem::solveSimplified(
@@ -2123,10 +2135,8 @@ bool ConstraintSystem::solveSimplified(
                                    allowFreeTypeVariables);
   }
 
-  if (disjunction) {
-    return !solveForDisjunctionChoices(disjunction, solutions,
-                                       allowFreeTypeVariables);
-  }
+  if (disjunction)
+    return solveForDisjunction(disjunction, solutions, allowFreeTypeVariables);
 
   // If there are no disjunctions we can't solve this system unless we have
   // free type variables and are allowing them in the solution.

lib/Sema/ConstraintSystem.h

@@ -2975,13 +2975,24 @@ class ConstraintSystem {
   /// \brief Collect the current inactive disjunction constraints.
   void collectDisjunctions(SmallVectorImpl<Constraint *> &disjunctions);
 
-  /// \brief Attempt to solve for the choices in the given disjunction.
-  ///
-  /// \returns true if we found at least one solution.
-  bool
-  solveForDisjunctionChoices(Constraint *disjunction,
-                             SmallVectorImpl<Solution> &solutions,
-                             FreeTypeVariableBinding allowFreeTypeVariables);
+  /// \brief Attempt find a solution involving the options in a
+  ///        disjunction.
+  ///
+  /// \returns true if we failed to find any solutions, false otherwise.
+  bool solveForDisjunction(Constraint *disjunction,
+                           SmallVectorImpl<Solution> &solutions,
+                           FreeTypeVariableBinding allowFreeTypeVariables);
+
+  /// \brief Attempt to solve for some subset of the constraints in a
+  ///        disjunction, skipping constraints that we decide do not
+  ///        need to be solved for because they would not result in
+  ///        the best solution to the constraint system.
+  ///
+  /// \returns true if we failed to find any solutions, false otherwise.
+  bool solveForDisjunctionChoices(
+      ArrayRef<Constraint *> constraints, ConstraintLocator *disjunctionLocator,
+      SmallVectorImpl<Solution> &solutions,
+      FreeTypeVariableBinding allowFreeTypeVariables, bool explicitConversion);
 
   /// \brief Solve the system of constraints after it has already been
   /// simplified.