[ConstraintGraph] Fix contractEdges to gather constraints only once

include/swift/Basic/Statistics.def

@@ -149,6 +149,11 @@ FRONTEND_STATISTIC(Sema, NumConformancesDeserialized)
 /// expression typechecker did".
 FRONTEND_STATISTIC(Sema, NumConstraintScopes)
 
+/// Number of constraints considered per attempt to
+/// contract constraint graph edges.
+/// This is a measure of complexity of the contraction algorithm.
+FRONTEND_STATISTIC(Sema, NumConstraintsConsideredForEdgeContraction)
+
 /// Number of declarations that were deserialized. A rough proxy for the amount
 /// of material loaded from other modules.
 FRONTEND_STATISTIC(Sema, NumDeclsDeserialized)

lib/Sema/ConstraintGraph.cpp

@@ -18,6 +18,7 @@
 #include "ConstraintGraph.h"
 #include "ConstraintGraphScope.h"
 #include "ConstraintSystem.h"
+#include "swift/Basic/Statistic.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include <algorithm>
@@ -27,6 +28,8 @@
 using namespace swift;
 using namespace constraints;
 
+#define DEBUG_TYPE "ConstraintGraph"
+
 #pragma mark Graph construction/destruction
 
 ConstraintGraph::ConstraintGraph(ConstraintSystem &cs) : CS(cs) { }
@@ -649,93 +652,89 @@ 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) {
+    // Track how many constraints did contraction algorithm iterated over.
+    incrementConstraintsPerContractionCounter();
+    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;
 }
 
@@ -773,6 +772,14 @@ void ConstraintGraph::optimize() {
   while (contractEdges()) {}
 }
 
+void ConstraintGraph::incrementConstraintsPerContractionCounter() {
+  SWIFT_FUNC_STAT;
+  auto &context = CS.getASTContext();
+  if (context.Stats)
+    context.Stats->getFrontendCounters()
+        .NumConstraintsConsideredForEdgeContraction++;
+}
+
 #pragma mark Debugging output
 
 void ConstraintGraphNode::print(llvm::raw_ostream &out, unsigned indent) {

lib/Sema/ConstraintGraph.h

@@ -332,6 +332,10 @@ class ConstraintGraph {
   /// Constraints that are "orphaned" because they contain no type variables.
   SmallVector<Constraint *, 4> OrphanedConstraints;
 
+  /// Increment the number of constraints considered per attempt
+  /// to contract constrant graph edges.
+  void incrementConstraintsPerContractionCounter();
+
   /// The kind of change made to the graph.
   enum class ChangeKind {
     /// Added a type variable.

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.

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

@@ -0,0 +1,10 @@
+// RUN: %scale-test --begin 0 --end 25000 --step 1000 --typecheck --select incrementConstraintsPerContractionCounter %s
+// REQUIRES: OS=macosx
+// REQUIRES: asserts
+
+let _: [Int] = [
+%for i in range(0, N):
+  1,
+%end
+  1
+]