[ConstraintSystem] Try to be more efficient solving linked operator expressions

lib/Sema/CSSolver.cpp

@@ -2086,9 +2086,72 @@ static Constraint *tryOptimizeGenericDisjunction(
   llvm_unreachable("covered switch");
 }
 
+// Performance hack: favor operator overloads with decl or type we're already
+// binding elsewhere in this expression.
+static void existingOperatorBindingsForDisjunction(ConstraintSystem &CS, ArrayRef<Constraint *> constraints, SmallVectorImpl<Constraint *> &found) {
+  auto *choice = constraints.front();
+  if (choice->getKind() != ConstraintKind::BindOverload)
+    return;
+
+  auto overload = choice->getOverloadChoice();
+  if (!overload.isDecl())
+    return;
+  auto decl = overload.getDecl();
+  if (!decl->isOperator())
+    return;
+  auto baseName = decl->getBaseName();
+
+  SmallSet<TypeBase *, 8> typesFound;
+
+  for (auto *resolved = CS.getResolvedOverloadSets(); resolved;
+       resolved = resolved->Previous) {
+    if (!resolved->Choice.isDecl())
+      continue;
+    auto representativeDecl = resolved->Choice.getDecl();
+
+    if (!representativeDecl->isOperator())
+      continue;
+
+    if (representativeDecl->getBaseName() == baseName) {
+      // Favor exactly the same decl, if we have a binding to the same name.
+      for (auto *constraint : constraints) {
+        if (constraint->isFavored())
+          continue;
+        auto choice = constraint->getOverloadChoice();
+        if (choice.getDecl() == representativeDecl) {
+          found.push_back(constraint);
+          break;
+        }
+      }
+    } else {
+      // Favor the same type, if we have a binding to an operator of that type.
+      auto representativeType = representativeDecl->getInterfaceType();
+      if (typesFound.count(representativeType.getPointer()))
+        continue;
+      typesFound.insert(representativeType.getPointer());
+
+      for (auto *constraint : constraints) {
+        if (constraint->isFavored())
+          continue;
+        auto choice = constraint->getOverloadChoice();
+        if (choice.getDecl()->getInterfaceType()->isEqual(representativeType)) {
+          found.push_back(constraint);
+          break;
+        }
+      }
+    }
+  }
+}
+
 void ConstraintSystem::partitionDisjunction(
     ArrayRef<Constraint *> Choices, SmallVectorImpl<unsigned> &Ordering,
     SmallVectorImpl<unsigned> &PartitionBeginning) {
+
+  SmallVector<Constraint *, 4> existing;
+  existingOperatorBindingsForDisjunction(*this, Choices, existing);
+  for (auto constraint : existing)
+    favorConstraint(constraint);
+
   // Apply a special-case rule for favoring one generic function over
   // another.
   if (auto favored = tryOptimizeGenericDisjunction(TC, DC, Choices)) {
@@ -2203,6 +2266,45 @@ void ConstraintSystem::partitionDisjunction(
   assert(Ordering.size() == Choices.size());
 }
 
+static Constraint *selectLeafApplyDisjunction(ConstraintSystem &cs, SmallVectorImpl<Constraint *> &disjunctions) {
+  // Only do this work for relatively deep search trees.
+  if (disjunctions.size() < 6)
+    return nullptr;
+
+  SmallDenseMap<Expr *, Constraint *> disjunctionForFnExpr;
+  SmallDenseMap<Constraint *, SourceRange> sourceRangeForDisjunction;
+
+  for (auto *disjunction : disjunctions)
+    disjunctionForFnExpr[disjunction->getLocator()->getAnchor()] = disjunction;
+
+  for (auto &constraint : cs.getConstraints()) {
+    if (constraint.getKind() != ConstraintKind::ApplicableFunction)
+      continue;
+
+    auto applyExpr = dyn_cast<ApplyExpr>(constraint.getLocator()->getAnchor());
+    if (!applyExpr)
+      continue;
+
+    if (auto disjunction = disjunctionForFnExpr[applyExpr->getFn()])
+      sourceRangeForDisjunction[disjunction] = applyExpr->getSourceRange();
+  }
+
+  Constraint *result = nullptr;
+  SourceRange range;
+  for (auto test : sourceRangeForDisjunction) {
+    if (range.isValid()) {
+      auto testRange = range;
+      testRange.widen(test.getSecond());
+      if (testRange != range)
+        continue;
+    }
+
+    result = test.getFirst();
+    range = test.getSecond();
+  }
+  return result;
+}
+
 Constraint *ConstraintSystem::selectDisjunction() {
   SmallVector<Constraint *, 4> disjunctions;
 
@@ -2224,18 +2326,22 @@ Constraint *ConstraintSystem::selectDisjunction() {
   if (auto *disjunction = selectBestBindingDisjunction(*this, disjunctions))
     return disjunction;
 
-  // Pick the disjunction with the smallest number of active choices.
   auto minDisjunction =
       std::min_element(disjunctions.begin(), disjunctions.end(),
                        [&](Constraint *first, Constraint *second) -> bool {
                          return first->countActiveNestedConstraints() <
                                 second->countActiveNestedConstraints();
                        });
 
-  if (minDisjunction != disjunctions.end())
-    return *minDisjunction;
+  if (minDisjunction == disjunctions.end())
+    return nullptr;
 
-  return nullptr;
+  // If the minimum sized disjunction is an overload disjunction, then try to choose a leaf apply expr from the set of disjunctions.
+  // Choosing an overload disjunction for an apply that depends on other disjunctions as arguments will likely not allow short-circuiting.
+  if ((*minDisjunction)->getNestedConstraints().front()->getKind() == ConstraintKind::BindOverload)
+    if (auto apply = selectLeafApplyDisjunction(*this, disjunctions))
+      return apply;
+  return *minDisjunction;
 }
 
 bool DisjunctionChoice::attempt(ConstraintSystem &cs) const {

lib/Sema/CSStep.cpp

@@ -334,7 +334,45 @@ StepResult ComponentStep::take(bool prevFailed) {
     // Produce a type variable step.
     return suspend(
         llvm::make_unique<TypeVariableStep>(CS, *bestBindings, Solutions));
-  } else if (disjunction) {
+  }
+
+  // Before stepping into a new disjunction, see if we can short-circuit this part of the solution tree by attempting to show that the last opened typevar is impossible to bind. This prunes exponential behavior when there are multiple generic function/operator overloads in the expression.
+  if (disjunction && CS.OpenedTypes.size()) {
+    auto lastOpenedType = CS.OpenedTypes.back();
+    auto lastOpenedTyvar = CS.getRepresentative(lastOpenedType.second.back().second);
+    if (CS.TypeVariables.count(lastOpenedTyvar) && !CS.getFixedType(lastOpenedTyvar)) {
+      auto bindings = CS.getPotentialBindings(lastOpenedTyvar);
+      if (!bindings.FullyBound && !bindings.PotentiallyIncomplete && bindings.Bindings.size()) {
+        size_t failures = 0;
+        for (auto binding : bindings.Bindings) {
+          // We can't definitively rule out some subtype being possible here.
+          if (binding.Kind == ConstraintSystem::AllowedBindingKind::Subtypes)
+            break;
+
+          ConstraintSystem::SolverScope scope(CS);
+          if (CS.TC.getLangOpts().DebugConstraintSolver) {
+            auto &log = CS.getASTContext().TypeCheckerDebug->getStream();
+            log.indent(CS.solverState ? CS.solverState->depth * 2 : 2)
+            << "(examining for short-circuit)\n";
+          }
+          CS.addConstraint(ConstraintKind::Bind, lastOpenedTyvar, binding.BindingType, lastOpenedType.first);
+          if (CS.simplify())
+            failures++;
+        }
+        if (failures == bindings.Bindings.size()) {
+          if (CS.TC.getLangOpts().DebugConstraintSolver) {
+            auto &log = CS.getASTContext().TypeCheckerDebug->getStream();
+            log.indent(CS.solverState ? CS.solverState->depth * 2 : 2)
+              << "(short-circuit due to last opened type binding failure "
+              << lastOpenedTyvar->getString() << ")\n";
+          }
+          return done(/*isSuccess=*/false);
+        }
+      }
+    }
+  }
+
+  if (disjunction) {
     // Produce a disjunction step.
     return suspend(
         llvm::make_unique<DisjunctionStep>(CS, disjunction, Solutions));
@@ -636,8 +674,12 @@ bool DisjunctionStep::shortCircuitDisjunctionAt(
   auto &ctx = CS.getASTContext();
 
   // If the successfully applied constraint is favored, we'll consider that to
-  // be the "best".
-  if (lastSuccessfulChoice->isFavored() && !currentChoice->isFavored()) {
+  // be the "best".  If it was only temporarily favored because it matched other
+  // operator bindings, we can even short-circuit other favored constraints.
+  if (lastSuccessfulChoice->isFavored() &&
+      (!currentChoice->isFavored() ||
+       (CS.solverState->isTemporarilyFavored(lastSuccessfulChoice) &&
+        !CS.solverState->isTemporarilyFavored(currentChoice)))) {
 #if !defined(NDEBUG)
     if (lastSuccessfulChoice->getKind() == ConstraintKind::BindOverload) {
       auto overloadChoice = lastSuccessfulChoice->getOverloadChoice();
@@ -676,16 +718,47 @@ 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() &&
       !ctx.LangOpts.SolverEnableOperatorDesignatedTypes) {
-    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, 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.
+    auto lastTyvar =
+        lastSuccessfulChoice->getFirstType()->getAs<TypeVariableType>();
+    auto lastRep = CS.getRepresentative(lastTyvar);
+
+    for (auto overload : Solutions.back().overloadChoices) {
+      auto overloadChoice = overload.getSecond().choice;
+      if (!overloadChoice.isDecl() ||
+          overloadChoice.getDecl() !=
+              lastSuccessfulChoice->getOverloadChoice().getDecl())
+        continue;
+
+      auto choiceTyvar =
+          CS.getType(simplifyLocatorToAnchor(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/Constraint.cpp

@@ -274,15 +274,17 @@ void Constraint::print(llvm::raw_ostream &Out, SourceManager *sm) const {
       Locator->dump(sm, Out);
       Out << "]]";
     }
-    Out << ":";
+    Out << ":\n";
 
     interleave(getNestedConstraints(),
                [&](Constraint *constraint) {
                  if (constraint->isDisabled())
-                   Out << "[disabled] ";
+                   Out << ">  [disabled] ";
+                 else
+                   Out << ">             ";
                  constraint->print(Out, sm);
                },
-               [&] { Out << " or "; });
+               [&] { Out << "\n"; });
     return;
   }
 

lib/Sema/ConstraintSystem.h

@@ -1436,6 +1436,16 @@ class ConstraintSystem {
       favoredConstraints.push_back(constraint);
     }
 
+    /// Whether or not the given constraint is only favored during this scope.
+    bool isTemporarilyFavored(Constraint *constraint) {
+      assert(constraint->isFavored());
+
+      for (auto test : favoredConstraints)
+        if (test == constraint)
+          return true;
+      return false;
+    }
+
   private:
     /// The list of constraints that have been retired along the
     /// current path, this list is used in LIFO fashion when constraints

test/Constraints/sr10324.swift

@@ -0,0 +1,19 @@
+// RUN: %target-swift-frontend -typecheck -verify %s
+
+struct A {
+    static func * (lhs: A, rhs: A) -> B { return B() }
+    static func * (lhs: B, rhs: A) -> B { return B() }
+    static func * (lhs: A, rhs: B) -> B { return B() }
+}
+struct B {}
+
+let (x, y, z) = (A(), A(), A())
+
+let w = A() * A() * A()     // works
+
+// Should all work
+let a = x * y * z
+let b = x * (y * z)
+let c = (x * y) * z
+let d = x * (y * z as B)
+let e = (x * y as B) * z

validation-test/Sema/type_checker_perf/fast/rdar18360240.swift

@@ -0,0 +1,6 @@
+// RUN: %target-typecheck-verify-swift -solver-expression-time-threshold=1
+// REQUIRES: OS=macosx
+// REQUIRES: asserts
+
+let empty: [Int] = []
+let _ = empty + empty + empty + empty + empty + empty + empty + empty + empty + empty + empty + empty

validation-test/Sema/type_checker_perf/fast/rdar25866240.swift

@@ -0,0 +1,7 @@
+// RUN: %target-typecheck-verify-swift -solver-expression-time-threshold=1
+// REQUIRES: OS=macosx
+// REQUIRES: asserts
+
+func f(c1: [String], c2: [String], c3: [String], c4: [String], c5: [String], c6: [String], c7: [String], c8: [String], c9: [String], c10: [String]) {
+  _ = c1 + c2 + c3 + c4 + c5 + c6 + c7 + c8 + c9 + c10
+}

validation-test/Sema/type_checker_perf/fast/simd_add.gyb

@@ -1,4 +1,4 @@
-// RUN: %scale-test --begin 3 --end 7 --step 1 --select NumLeafScopes %s
+// RUN: %scale-test --begin 3 --end 7 --step 1 --select NumLeafScopes %s --exponential-threshold 11.0
 // REQUIRES: OS=macosx
 // REQUIRES: asserts