[ConstraintGraph] Fix contractEdges to gather constraints only once

Currently we have this non-optimal behavior in `contractEdges` where
for every type variable it gathers constraints for its whole equivalence
class before checking if any are "contractable", instead constraints
could be gathered/filtered once which removes a lot of useless work.

Author: xedin

lib/Sema/ConstraintGraph.cpp

@@ -649,93 +649,87 @@ static bool shouldContractEdge(ConstraintKind kind) {
 }
 
 bool ConstraintGraph::contractEdges() {
-  llvm::SetVector<std::pair<TypeVariableType *,
-                            TypeVariableType *>> contractions;
-
-  auto tyvars = getTypeVariables();
-  auto didContractEdges = false;
+  SmallVector<Constraint *, 16> constraints;
+  CS.findConstraints(constraints, [](const Constraint &constraint) {
+    return shouldContractEdge(constraint.getKind());
+  });
 
-  for (auto tyvar : tyvars) {
-    SmallVector<Constraint *, 4> constraints;
-    gatherConstraints(tyvar, constraints,
-                      ConstraintGraph::GatheringKind::EquivalenceClass);
+  bool didContractEdges = false;
+  for (auto *constraint : constraints) {
+    auto kind = constraint->getKind();
 
-    for (auto constraint : constraints) {
-      auto kind = constraint->getKind();
-      // Contract binding edges between type variables.
-      if (shouldContractEdge(kind)) {
-        auto t1 = constraint->getFirstType()->getDesugaredType();
-        auto t2 = constraint->getSecondType()->getDesugaredType();
+    // Contract binding edges between type variables.
+    assert(shouldContractEdge(kind));
 
-        auto tyvar1 = t1->getAs<TypeVariableType>();
-        auto tyvar2 = t2->getAs<TypeVariableType>();
+    auto t1 = constraint->getFirstType()->getDesugaredType();
+    auto t2 = constraint->getSecondType()->getDesugaredType();
 
-        if (!(tyvar1 && tyvar2))
-          continue;
+    auto tyvar1 = t1->getAs<TypeVariableType>();
+    auto tyvar2 = t2->getAs<TypeVariableType>();
 
-        auto isParamBindingConstraint = kind == ConstraintKind::BindParam;
-
-        // If the argument is allowed to bind to `inout`, in general,
-        // it's invalid to contract the edge between argument and parameter,
-        // but if we can prove that there are no possible bindings
-        // which result in attempt to bind `inout` type to argument
-        // type variable, we should go ahead and allow (temporary)
-        // contraction, because that greatly helps with performance.
-        // Such action is valid because argument type variable can
-        // only get its bindings from related overload, which gives
-        // us enough information to decided on l-valueness.
-        if (isParamBindingConstraint && tyvar1->getImpl().canBindToInOut()) {
-          bool isNotContractable = true;
-          if (auto bindings = CS.getPotentialBindings(tyvar1)) {
-            for (auto &binding : bindings.Bindings) {
-              auto type = binding.BindingType;
-              isNotContractable = type.findIf([&](Type nestedType) -> bool {
-                if (auto tv = nestedType->getAs<TypeVariableType>()) {
-                  if (!tv->getImpl().mustBeMaterializable())
-                    return true;
-                }
-
-                return nestedType->is<InOutType>();
-              });
+    if (!(tyvar1 && tyvar2))
+      continue;
 
-              // If there is at least one non-contractable binding, let's
-              // not risk contracting this edge.
-              if (isNotContractable)
-                break;
+    auto isParamBindingConstraint = kind == ConstraintKind::BindParam;
+
+    // If the argument is allowed to bind to `inout`, in general,
+    // it's invalid to contract the edge between argument and parameter,
+    // but if we can prove that there are no possible bindings
+    // which result in attempt to bind `inout` type to argument
+    // type variable, we should go ahead and allow (temporary)
+    // contraction, because that greatly helps with performance.
+    // Such action is valid because argument type variable can
+    // only get its bindings from related overload, which gives
+    // us enough information to decided on l-valueness.
+    if (isParamBindingConstraint && tyvar1->getImpl().canBindToInOut()) {
+      bool isNotContractable = true;
+      if (auto bindings = CS.getPotentialBindings(tyvar1)) {
+        for (auto &binding : bindings.Bindings) {
+          auto type = binding.BindingType;
+          isNotContractable = type.findIf([&](Type nestedType) -> bool {
+            if (auto tv = nestedType->getAs<TypeVariableType>()) {
+              if (!tv->getImpl().mustBeMaterializable())
+                return true;
             }
-          }
 
+            return nestedType->is<InOutType>();
+          });
+
+          // If there is at least one non-contractable binding, let's
+          // not risk contracting this edge.
           if (isNotContractable)
-            continue;
+            break;
         }
+      }
 
-        auto rep1 = CS.getRepresentative(tyvar1);
-        auto rep2 = CS.getRepresentative(tyvar2);
-
-        if (((rep1->getImpl().canBindToLValue() ==
-              rep2->getImpl().canBindToLValue()) ||
-              // Allow l-value contractions when binding parameter types.
-              isParamBindingConstraint)) {
-          if (CS.TC.getLangOpts().DebugConstraintSolver) {
-            auto &log = CS.getASTContext().TypeCheckerDebug->getStream();
-            if (CS.solverState)
-              log.indent(CS.solverState->depth * 2);
-
-            log << "Contracting constraint ";
-            constraint->print(log, &CS.getASTContext().SourceMgr);
-            log << "\n";
-          }
-
-          // Merge the edges and remove the constraint.
-          removeEdge(constraint);
-          if (rep1 != rep2)
-            CS.mergeEquivalenceClasses(rep1, rep2, /*updateWorkList*/ false);
-          didContractEdges = true;
-        }
+      if (isNotContractable)
+        continue;
+    }
+
+    auto rep1 = CS.getRepresentative(tyvar1);
+    auto rep2 = CS.getRepresentative(tyvar2);
+
+    if (((rep1->getImpl().canBindToLValue() ==
+          rep2->getImpl().canBindToLValue()) ||
+         // Allow l-value contractions when binding parameter types.
+         isParamBindingConstraint)) {
+      if (CS.TC.getLangOpts().DebugConstraintSolver) {
+        auto &log = CS.getASTContext().TypeCheckerDebug->getStream();
+        if (CS.solverState)
+          log.indent(CS.solverState->depth * 2);
+
+        log << "Contracting constraint ";
+        constraint->print(log, &CS.getASTContext().SourceMgr);
+        log << "\n";
       }
+
+      // Merge the edges and remove the constraint.
+      removeEdge(constraint);
+      if (rep1 != rep2)
+        CS.mergeEquivalenceClasses(rep1, rep2, /*updateWorkList*/ false);
+      didContractEdges = true;
     }
   }
-
   return didContractEdges;
 }
 

lib/Sema/ConstraintSystem.h

@@ -1909,6 +1909,12 @@ class ConstraintSystem {
   /// due to a change.
   ConstraintList &getActiveConstraints() { return ActiveConstraints; }
 
+  void findConstraints(SmallVectorImpl<Constraint *> &found,
+                       llvm::function_ref<bool(const Constraint &)> pred) {
+    filterConstraints(ActiveConstraints, pred, found);
+    filterConstraints(InactiveConstraints, pred, found);
+  }
+
   /// \brief Retrieve the representative of the equivalence class containing
   /// this type variable.
   TypeVariableType *getRepresentative(TypeVariableType *typeVar) {
@@ -2074,6 +2080,16 @@ class ConstraintSystem {
   /// into the worklist.
   void addTypeVariableConstraintsToWorkList(TypeVariableType *typeVar);
 
+  static void
+  filterConstraints(ConstraintList &constraints,
+                    llvm::function_ref<bool(const Constraint &)> pred,
+                    SmallVectorImpl<Constraint *> &found) {
+    for (auto &constraint : constraints) {
+      if (pred(constraint))
+        found.push_back(&constraint);
+    }
+  }
+
 public:
 
   /// \brief Coerce the given expression to an rvalue, if it isn't already.