Instead of chaining binops, favor disjunctions with op overloads whose types match existing binding choices

gregomni · hborla · commit 2edba9dfbd88 · 2020-10-15T09:03:45.000-07:00
diff --git a/include/swift/Sema/Constraint.h b/include/swift/Sema/Constraint.h
@@ -672,6 +672,15 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return Nested;
   }
 
+  unsigned countFavoredNestedConstraints() const {
+    unsigned count = 0;
+    for (auto *constraint : Nested)
+      if (constraint->isFavored() && !constraint->isDisabled())
+        count++;
+
+    return count;
+  }
+
   unsigned countActiveNestedConstraints() const {
     unsigned count = 0;
     for (auto *constraint : Nested)
diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
@@ -46,26 +46,6 @@ static bool isArithmeticOperatorDecl(ValueDecl *vd) {
    vd->getBaseName() == "%");
 }
 
-static bool hasBinOpOverloadWithSameArgTypesDifferingResult(
-    OverloadedDeclRefExpr *overloads) {
-  for (auto decl : overloads->getDecls()) {
-    auto metaFuncType = decl->getInterfaceType()->getAs<AnyFunctionType>();
-    auto funcType = metaFuncType->getResult()->getAs<FunctionType>();
-    if (!funcType)
-      continue;
-
-    auto params = funcType->getParams();
-    if (params.size() != 2)
-      continue;
-
-    if (params[0].getPlainType().getPointer() != params[1].getPlainType().getPointer())
-      continue;
-    if (funcType->getResult().getPointer() != params[0].getPlainType().getPointer())
-      return true;
-  }
-  return false;
-}
-
 static bool mergeRepresentativeEquivalenceClasses(ConstraintSystem &CS,
                                                   TypeVariableType* tyvar1,
                                                   TypeVariableType* tyvar2) {
@@ -323,75 +303,6 @@ namespace {
       // solver can still make progress.
       auto favoredTy = (*lti.collectedTypes.begin())->getWithoutSpecifierType();
       CS.setFavoredType(expr, favoredTy.getPointer());
-
-      // If we have a chain of identical binop expressions with homogeneous
-      // argument types, we can directly simplify the associated constraint
-      // graph.
-      auto simplifyBinOpExprTyVars = [&]() {
-        // Don't attempt to do linking if there are
-        // literals intermingled with other inferred types.
-        if (lti.hasLiteral)
-          return;
-
-        for (auto binExp1 : lti.binaryExprs) {
-          for (auto binExp2 : lti.binaryExprs) {
-            if (binExp1 == binExp2)
-              continue;
-
-            auto fnTy1 = CS.getType(binExp1)->getAs<TypeVariableType>();
-            auto fnTy2 = CS.getType(binExp2)->getAs<TypeVariableType>();
-
-            if (!(fnTy1 && fnTy2))
-              return;
-
-            auto ODR1 = dyn_cast<OverloadedDeclRefExpr>(binExp1->getFn());
-            auto ODR2 = dyn_cast<OverloadedDeclRefExpr>(binExp2->getFn());
-
-            if (!(ODR1 && ODR2))
-              return;
-
-            // TODO: We currently limit this optimization to known arithmetic
-            // operators, but we should be able to broaden this out to
-            // logical operators as well.
-            if (!isArithmeticOperatorDecl(ODR1->getDecls()[0]))
-              return;
-
-            if (ODR1->getDecls()[0]->getBaseName() !=
-                ODR2->getDecls()[0]->getBaseName())
-              return;
-
-            if (hasBinOpOverloadWithSameArgTypesDifferingResult(ODR1))
-              return;
-
-            // All things equal, we can merge the tyvars for the function
-            // types.
-            auto rep1 = CS.getRepresentative(fnTy1);
-            auto rep2 = CS.getRepresentative(fnTy2);
-
-            if (rep1 != rep2) {
-              CS.mergeEquivalenceClasses(rep1, rep2,
-                                         /*updateWorkList*/ false);
-            }
-
-            auto odTy1 = CS.getType(ODR1)->getAs<TypeVariableType>();
-            auto odTy2 = CS.getType(ODR2)->getAs<TypeVariableType>();
-
-            if (odTy1 && odTy2) {
-              auto odRep1 = CS.getRepresentative(odTy1);
-              auto odRep2 = CS.getRepresentative(odTy2);
-
-              // Since we'll be choosing the same overload, we can merge
-              // the overload tyvar as well.
-              if (odRep1 != odRep2)
-                CS.mergeEquivalenceClasses(odRep1, odRep2,
-                                           /*updateWorkList*/ false);
-            }
-          }
-        }
-      };
-
-      simplifyBinOpExprTyVars();       
-
       return true;
     }
     
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -2013,9 +2013,47 @@ static Constraint *tryOptimizeGenericDisjunction(
   llvm_unreachable("covered switch");
 }
 
+// Performance hack: favor operator overloads with types 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;
+
+  for (auto *resolved = CS.getResolvedOverloadSets(); resolved;
+       resolved = resolved->Previous) {
+    if (!resolved->Choice.isDecl())
+      continue;
+    auto representativeDecl = resolved->Choice.getDecl();
+
+    if (!representativeDecl->isOperator())
+      continue;
+
+    for (auto *constraint : constraints) {
+      if (constraint->isFavored())
+        continue;
+      auto choice = constraint->getOverloadChoice();
+      if (choice.getDecl()->getInterfaceType()->isEqual(representativeDecl->getInterfaceType()))
+        found.push_back(constraint);
+    }
+  }
+}
+
 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(DC, Choices)) {
@@ -2134,12 +2172,30 @@ Constraint *ConstraintSystem::selectDisjunction() {
   if (auto *disjunction = selectBestBindingDisjunction(*this, disjunctions))
     return disjunction;
 
-  // Pick the disjunction with the smallest number of active choices.
+  // 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 {
-                         return first->countActiveNestedConstraints() <
-                                second->countActiveNestedConstraints();
+                         unsigned firstFavored = first->countFavoredNestedConstraints();
+                         unsigned secondFavored = second->countFavoredNestedConstraints();
+
+                         if (firstFavored == secondFavored) {
+                           // Look for additional choices to favor
+                           SmallVector<Constraint *, 4> firstExisting;
+                           SmallVector<Constraint *, 4> secondExisting;
+
+                           existingOperatorBindingsForDisjunction(*cs, first->getNestedConstraints(), firstExisting);
+                           firstFavored = firstExisting.size() ?: first->countActiveNestedConstraints();
+                           existingOperatorBindingsForDisjunction(*cs, second->getNestedConstraints(), secondExisting);
+                           secondFavored = secondExisting.size() ?: second->countActiveNestedConstraints();
+
+                           return firstFavored < secondFavored;
+                         } else {
+                           firstFavored = firstFavored ?: first->countActiveNestedConstraints();
+                           secondFavored = secondFavored ?: second->countActiveNestedConstraints();
+                           return firstFavored < secondFavored;
+                         }
                        });
 
   if (minDisjunction != disjunctions.end())
diff --git a/lib/Sema/CSStep.h b/lib/Sema/CSStep.h
@@ -654,7 +654,6 @@ 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;
   }
 
@@ -702,43 +701,6 @@ 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 elt : getResolvedOverloads()) {
-      auto resolved = elt.second;
-      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());
