Merge pull request #34399 from hborla/optimize-linked-operator-solving

[Constraint System] Implement heuristics for linked operator expressions in the solver proper.

Author: hborla

include/swift/Sema/Constraint.h

@@ -672,13 +672,16 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return Nested;
   }
 
-  unsigned countActiveNestedConstraints() const {
-    unsigned count = 0;
-    for (auto *constraint : Nested)
-      if (!constraint->isDisabled())
-        count++;
+  unsigned countFavoredNestedConstraints() const {
+    return llvm::count_if(Nested, [](const Constraint *constraint) {
+      return constraint->isFavored() && !constraint->isDisabled();
+    });
+  }
 
-    return count;
+  unsigned countActiveNestedConstraints() const {
+    return llvm::count_if(Nested, [](const Constraint *constraint) {
+      return !constraint->isDisabled();
+    });
   }
 
   /// Determine if this constraint represents explicit conversion,

include/swift/Sema/ConstraintSystem.h

@@ -5328,6 +5328,12 @@ class ConstraintSystem {
                             SmallVectorImpl<unsigned> &Ordering,
                             SmallVectorImpl<unsigned> &PartitionBeginning);
 
+  /// The overload sets that have already been resolved along the current path.
+  const llvm::MapVector<ConstraintLocator *, SelectedOverload> &
+  getResolvedOverloads() const {
+    return ResolvedOverloads;
+  }
+
   /// If we aren't certain that we've emitted a diagnostic, emit a fallback
   /// diagnostic.
   void maybeProduceFallbackDiagnostic(SolutionApplicationTarget target) const;

lib/Sema/CSGen.cpp

@@ -367,7 +367,7 @@ namespace {
         }
       };
 
-      simplifyBinOpExprTyVars();       
+      simplifyBinOpExprTyVars();
 
       return true;
     }

lib/Sema/CSSolver.cpp

@@ -2013,6 +2013,66 @@ static Constraint *tryOptimizeGenericDisjunction(
   llvm_unreachable("covered switch");
 }
 
+/// Populates the \c found vector with the indices of the given constraints
+/// that have a matching type to an existing operator binding elsewhere in
+/// the expression.
+///
+/// Operator bindings that have a matching type to an existing binding
+/// are attempted first by the solver because it's very common to chain
+/// operators of the same type together.
+static void existingOperatorBindingsForDisjunction(ConstraintSystem &CS,
+                                                   ArrayRef<Constraint *> constraints,
+                                                   SmallVectorImpl<unsigned> &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;
+
+  // For concrete operators, consider overloads that have the same type as
+  // an existing binding, because it's very common to write mixed operator
+  // expressions where all operands have the same type, e.g. `(x + 10) / 2`.
+  // For generic operators, only favor an exact overload that has already
+  // been bound, because mixed operator expressions are far less common, and
+  // computing generic canonical types is expensive.
+  SmallSet<CanType, 4> concreteTypesFound;
+  SmallSet<ValueDecl *, 4> genericDeclsFound;
+  for (auto overload : CS.getResolvedOverloads()) {
+    auto resolved = overload.second;
+    if (!resolved.choice.isDecl())
+      continue;
+
+    auto representativeDecl = resolved.choice.getDecl();
+    if (!representativeDecl->isOperator())
+      continue;
+
+    auto interfaceType = representativeDecl->getInterfaceType();
+    if (interfaceType->is<GenericFunctionType>()) {
+      genericDeclsFound.insert(representativeDecl);
+    } else {
+      concreteTypesFound.insert(interfaceType->getCanonicalType());
+    }
+  }
+
+  for (auto index : indices(constraints)) {
+    auto *constraint = constraints[index];
+    if (constraint->isFavored())
+      continue;
+
+    auto *decl = constraint->getOverloadChoice().getDecl();
+    auto interfaceType = decl->getInterfaceType();
+    bool isGeneric = interfaceType->is<GenericFunctionType>();
+    if ((isGeneric && genericDeclsFound.count(decl)) ||
+        (!isGeneric && concreteTypesFound.count(interfaceType->getCanonicalType())))
+      found.push_back(index);
+  }
+}
+
 void ConstraintSystem::partitionDisjunction(
     ArrayRef<Constraint *> Choices, SmallVectorImpl<unsigned> &Ordering,
     SmallVectorImpl<unsigned> &PartitionBeginning) {
@@ -2042,12 +2102,18 @@ void ConstraintSystem::partitionDisjunction(
 
   // First collect some things that we'll generally put near the beginning or
   // end of the partitioning.
-
   SmallVector<unsigned, 4> favored;
+  SmallVector<unsigned, 4> everythingElse;
   SmallVector<unsigned, 4> simdOperators;
   SmallVector<unsigned, 4> disabled;
   SmallVector<unsigned, 4> unavailable;
 
+  // Add existing operator bindings to the main partition first. This often
+  // helps the solver find a solution fast.
+  existingOperatorBindingsForDisjunction(*this, Choices, everythingElse);
+  for (auto index : everythingElse)
+    taken.insert(Choices[index]);
+
   // First collect disabled and favored constraints.
   forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
     if (constraint->isDisabled()) {
@@ -2107,7 +2173,6 @@ void ConstraintSystem::partitionDisjunction(
         }
       };
 
-  SmallVector<unsigned, 4> everythingElse;
   // Gather the remaining options.
   forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
     everythingElse.push_back(index);
@@ -2134,13 +2199,34 @@ 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();
-                       });
+  // Pick the disjunction with the smallest number of favored, then active choices.
+  auto cs = this;
+  auto minDisjunction = std::min_element(disjunctions.begin(), disjunctions.end(),
+      [&](Constraint *first, Constraint *second) -> bool {
+        unsigned firstFavored = first->countFavoredNestedConstraints();
+        unsigned secondFavored = second->countFavoredNestedConstraints();
+
+        if (!isOperatorBindOverload(first->getNestedConstraints().front()) ||
+            !isOperatorBindOverload(second->getNestedConstraints().front()))
+          return first->countActiveNestedConstraints() < second->countActiveNestedConstraints();
+
+        if (firstFavored == secondFavored) {
+          // Look for additional choices to favor
+          SmallVector<unsigned, 4> firstExisting;
+          SmallVector<unsigned, 4> secondExisting;
+
+          existingOperatorBindingsForDisjunction(*cs, first->getNestedConstraints(), firstExisting);
+          firstFavored = firstExisting.size() ? firstExisting.size() : first->countActiveNestedConstraints();
+          existingOperatorBindingsForDisjunction(*cs, second->getNestedConstraints(), secondExisting);
+          secondFavored = secondExisting.size() ? secondExisting.size() : second->countActiveNestedConstraints();
+
+          return firstFavored < secondFavored;
+        }
+
+        firstFavored = firstFavored ? firstFavored : first->countActiveNestedConstraints();
+        secondFavored = secondFavored ? secondFavored : second->countActiveNestedConstraints();
+        return firstFavored < secondFavored;
+      });
 
   if (minDisjunction != disjunctions.end())
     return *minDisjunction;

lib/Sema/CSStep.cpp

@@ -616,21 +616,6 @@ bool DisjunctionStep::shortCircuitDisjunctionAt(
     Constraint *currentChoice, Constraint *lastSuccessfulChoice) const {
   auto &ctx = CS.getASTContext();
 
-  // If the successfully applied constraint is favored, we'll consider that to
-  // be the "best".
-  if (lastSuccessfulChoice->isFavored() && !currentChoice->isFavored()) {
-#if !defined(NDEBUG)
-    if (lastSuccessfulChoice->getKind() == ConstraintKind::BindOverload) {
-      auto overloadChoice = lastSuccessfulChoice->getOverloadChoice();
-      assert((!overloadChoice.isDecl() ||
-              !overloadChoice.getDecl()->getAttrs().isUnavailable(ctx)) &&
-             "Unavailable decl should not be favored!");
-    }
-#endif
-
-    return true;
-  }
-
   // Anything without a fix is better than anything with a fix.
   if (currentChoice->getFix() && !lastSuccessfulChoice->getFix())
     return true;
@@ -657,15 +642,6 @@ 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.
-  if (currentChoice->getKind() == ConstraintKind::BindOverload &&
-      isSIMDOperator(currentChoice->getOverloadChoice().getDecl()) &&
-      lastSuccessfulChoice->getKind() == ConstraintKind::BindOverload &&
-      !isSIMDOperator(lastSuccessfulChoice->getOverloadChoice().getDecl())) {
-    return true;
-  }
-
   return false;
 }
 

lib/Sema/CSStep.h

@@ -213,11 +213,6 @@ class SolverStep {
     CS.CG.addConstraint(constraint);
   }
 
-  const llvm::MapVector<ConstraintLocator *, SelectedOverload> &
-  getResolvedOverloads() const {
-    return CS.ResolvedOverloads;
-  }
-
   void recordDisjunctionChoice(ConstraintLocator *disjunctionLocator,
                                unsigned index) const {
     CS.recordDisjunctionChoice(disjunctionLocator, index);
@@ -716,8 +711,8 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
     if (!repr || repr == typeVar)
       return;
 
-    for (auto elt : getResolvedOverloads()) {
-      auto resolved = elt.second;
+    for (auto overload : CS.getResolvedOverloads()) {
+      auto resolved = overload.second;
       if (!resolved.boundType->isEqual(repr))
         continue;
 

lib/Sema/Constraint.cpp

@@ -317,15 +317,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;
   }
 

test/Constraints/sr10324.swift

@@ -0,0 +1,21 @@
+// RUN: %target-swift-frontend -typecheck -verify %s
+
+// REQUIRES: rdar65007946
+
+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/slow/expression_too_complex_4.swift renamed to validation-test/Sema/type_checker_perf/fast/expression_too_complex_4.swift

@@ -5,5 +5,4 @@ func test(_ i: Int, _ j: Int) -> Int {
   return 1 + (((i >> 1) + (i >> 2) + (i >> 3) + (i >> 4) << 1) << 1) & 0x40 +
          1 + (((i >> 1) + (i >> 2) + (i >> 3) + (i >> 4) << 1) << 1) & 0x40 +
          1 + (((i >> 1) + (i >> 2) + (i >> 3) + (i >> 4) << 1) << 1) & 0x40
-  // expected-error@-1 {{the compiler is unable to type-check this expression in reasonable time}}
 }

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

@@ -1,4 +1,4 @@
-// RUN: %scale-test --begin 2 --end 10 --step 2 --select NumConstraintScopes --polynomial-threshold 1.5 %s
+// RUN: %scale-test --begin 2 --end 10 --step 2 --select NumConstraintScopes %s
 // REQUIRES: asserts,no_asan
 
 let empty: [Int] = []

validation-test/Sema/type_checker_perf/slow/rdar22022980.swift renamed to validation-test/Sema/type_checker_perf/fast/rdar22022980.swift

@@ -2,4 +2,3 @@
 // REQUIRES: tools-release,no_asan
 
 _ = [1, 3, 5, 7, 11].filter{ $0 == 1 || $0 == 3 || $0 == 11 || $0 == 1 || $0 == 3 || $0 == 11 } == [ 1, 3, 11 ]
-// expected-error@-1 {{unable to type-check}}

validation-test/Sema/type_checker_perf/slow/rdar23429943.swift renamed to validation-test/Sema/type_checker_perf/fast/rdar23429943.swift

@@ -1,7 +1,6 @@
 // RUN: %target-typecheck-verify-swift -solver-expression-time-threshold=1
 // REQUIRES: tools-release,no_asan
 
-// expected-error@+1 {{the compiler is unable to type-check this expression in reasonable time}}
 let _ = [0].reduce([Int]()) {
   return $0.count == 0 && ($1 == 0 || $1 == 2 || $1 == 3) ? [] : $0 + [$1]
 }

validation-test/Sema/type_checker_perf/slow/rdar23861629.swift renamed to validation-test/Sema/type_checker_perf/fast/rdar23861629.swift

@@ -5,7 +5,6 @@ struct S { var s: String? }
 
 func rdar23861629(_ a: [S]) {
   _ = a.reduce("") {
-    // expected-error@-1 {{reasonable time}}
     ($0 == "") ? ($1.s ?? "") : ($0 + "," + ($1.s ?? "")) + ($1.s ?? "test") + ($1.s ?? "okay")
   }
 }

validation-test/Sema/type_checker_perf/slow/mixed_string_array_addition.swift

@@ -0,0 +1,12 @@
+// RUN: %target-typecheck-verify-swift -swift-version 5 -solver-expression-time-threshold=1
+
+func method(_ arg: String, body: () -> [String]) {}
+
+func test(str: String, properties: [String]) {
+  // expected-error@+1 {{the compiler is unable to type-check this expression in reasonable time}}
+  method(str + "" + str + "") {
+    properties.map { param in
+      "" + param + "" + param + ""
+    } + [""]
+  }
+}