Merge typevars of operators with shared decls after finding a solution to speed up further searching

Author: gregomni

lib/Sema/CSStep.cpp

@@ -659,15 +659,50 @@ bool DisjunctionStep::shortCircuitDisjunctionAt(
   if (currentChoice->getKind() == ConstraintKind::CheckedCast)
     return true;
 
-  // If we have a SIMD operator, and the prior choice was not a SIMD
-  // Operator, we're done.
+  // Extra checks for binding of operators
   if (currentChoice->getKind() == ConstraintKind::BindOverload &&
-      isSIMDOperator(currentChoice->getOverloadChoice().getDecl()) &&
+      currentChoice->getOverloadChoice().getDecl()->isOperator() &&
       lastSuccessfulChoice->getKind() == ConstraintKind::BindOverload &&
-      !isSIMDOperator(lastSuccessfulChoice->getOverloadChoice().getDecl())) {
+      lastSuccessfulChoice->getOverloadChoice().getDecl()->isOperator()) {
     return true;
   }
 
+    // If we have a SIMD operator, and the prior choice was not a SIMD
+    // Operator, we're done.
+    if (isSIMDOperator(currentChoice->getOverloadChoice().getDecl()) &&
+        !isSIMDOperator(lastSuccessfulChoice->getOverloadChoice().getDecl()))
+      return true;
+
+    // Otherwise if we have an existing solution, bind tyvars bound to the same
+    // decl in the solution to the choice tyvar. We can continue finding more
+    // solutions, but all the instances of the operator that chose the same
+    // overload as this successful choice will be bound togeter.
+    if (Solutions.size()) {
+      auto lastTyvar =
+          lastSuccessfulChoice->getFirstType()->getAs<TypeVariableType>();
+      auto lastRep = CS.getRepresentative(lastTyvar);
+
+      for (auto overload : Solutions[0].overloadChoices) {
+        auto overloadChoice = overload.getSecond().choice;
+        if (!overloadChoice.isDecl() ||
+            overloadChoice.getDecl() !=
+                lastSuccessfulChoice->getOverloadChoice().getDecl())
+          continue;
+
+        auto choiceTyvar =
+            CS.getType(simplifyLocatorToAnchor(CS, overload.getFirst()))
+                ->getAs<TypeVariableType>();
+        if (!choiceTyvar)
+          continue;
+
+        auto rep = CS.getRepresentative(choiceTyvar);
+        if (lastRep != rep) {
+          CS.mergeEquivalenceClasses(rep, lastRep);
+          lastRep = CS.getRepresentative(lastRep);
+        }
+      }
+    }
+  }
   return false;
 }
 

lib/Sema/CSStep.h

@@ -649,6 +649,7 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
       : BindingStep(cs, {cs, disjunction}, solutions), Disjunction(disjunction),
         AfterDisjunction(erase(disjunction)) {
     assert(Disjunction->getKind() == ConstraintKind::Disjunction);
+    pruneOverloadSet(Disjunction);
     ++cs.solverState->NumDisjunctions;
   }
 
@@ -696,6 +697,43 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
   /// simplified further, false otherwise.
   bool attempt(const DisjunctionChoice &choice) override;
 
+  // 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.
+  void pruneOverloadSet(Constraint *disjunction) {
+    auto *choice = disjunction->getNestedConstraints().front();
+    auto *typeVar = choice->getFirstType()->getAs<TypeVariableType>();
+    if (!typeVar)
+      return;
+
+    auto *repr = typeVar->getImpl().getRepresentative(nullptr);
+    if (!repr || repr == typeVar)
+      return;
+
+    for (auto *resolved = getResolvedOverloads(); resolved;
+         resolved = resolved->Previous) {
+      if (!resolved->BoundType->isEqual(repr))
+        continue;
+
+      auto &representative = resolved->Choice;
+      if (!representative.isDecl())
+        return;
+
+      // 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() || choice.getDecl() == representative.getDecl())
+          continue;
+
+        constraint->setDisabled();
+        DisabledChoices.push_back(constraint);
+      }
+      break;
+    }
+  };
+
   // Figure out which of the solutions has the smallest score.
   static Optional<Score> getBestScore(SmallVectorImpl<Solution> &solutions) {
     assert(!solutions.empty());