Merge pull request swiftlang#76962 from slavapestov/connected-components-graph

Sema: Optimize computeConnectedComponents()

Author: slavapestov

lib/Sema/ConstraintGraph.cpp

@@ -722,6 +722,10 @@ namespace {
     mutable llvm::SmallDenseMap<TypeVariableType *, TypeVariableType *>
         representatives;
 
+    // Figure out which components have unbound type variables and/or
+    // constraints. These are the only components we want to report.
+    llvm::SmallDenseSet<TypeVariableType *> validComponents;
+
     /// The complete set of constraints that were visited while computing
     /// connected components.
     llvm::SmallPtrSet<Constraint *, 8> visitedConstraints;
@@ -764,47 +768,34 @@ namespace {
 
     /// Retrieve the set of components.
     SmallVector<Component, 1> getComponents() const {
-      // Figure out which components have unbound type variables and/or
-      // constraints. These are the only components we want to report.
-      llvm::SmallDenseSet<TypeVariableType *> validComponents;
-      auto &cs = cg.getConstraintSystem();
-      for (auto typeVar : typeVars) {
-        // If this type variable has a fixed type, skip it.
-        if (cs.getFixedType(typeVar))
-          continue;
+      // The final return value.
+      SmallVector<Component, 1> flatComponents;
 
-        auto rep = findRepresentative(typeVar);
-        validComponents.insert(rep);
-      }
 
-      for (auto &constraint : cs.getConstraints()) {
-        for (auto typeVar : constraint.getTypeVariables()) {
-          auto rep = findRepresentative(typeVar);
-          validComponents.insert(rep);
-        }
-      }
+      // We don't actually need to partition the graph into components if
+      // there are fewer than 2.
+      if (validComponents.size() < 2 && cg.getOrphanedConstraints().empty())
+        return flatComponents;
 
-      // Capture the type variables of each component.
+      // Mapping from representatives to components.
       llvm::SmallDenseMap<TypeVariableType *, Component> components;
       SmallVector<TypeVariableType *, 4> representativeTypeVars;
+
+      // Capture the type variables of each component.
       for (auto typeVar : typeVars) {
         // Find the representative. If we aren't creating a type variable
         // for this component, skip it.
         auto rep = findRepresentative(typeVar);
         if (validComponents.count(rep) == 0)
           continue;
 
-        // If this type variable is the representative, add it to the list of
-        // representatives.
-        if (rep == typeVar) {
+        auto pair = components.insert({rep, Component(components.size())});
+        if (pair.second)
           representativeTypeVars.push_back(rep);
-        }
 
         // Record this type variable in the set of type variables for its
         // component.
-        auto &component = components.insert(
-            {rep, Component(components.size())}).first->second;
-        component.typeVars.push_back(typeVar);
+        pair.first->second.typeVars.push_back(typeVar);
       }
 
       // Retrieve the component for the given representative type variable.
@@ -814,6 +805,8 @@ namespace {
         return component->second;
       };
 
+      auto &cs = cg.getConstraintSystem();
+
       // Assign each constraint to its appropriate component.
       // Note: we use the inactive constraints so that we maintain the
       // order of constraints when we re-introduce them.
@@ -848,8 +841,7 @@ namespace {
                   });
 
         representativeTypeVars =
-            computeOneWayComponentOrdering(representativeTypeVars,
-                                           validComponents);
+            computeOneWayComponentOrdering(representativeTypeVars);
 
         // Fill in one-way dependency information for all of the components.
         for (auto typeVar : representativeTypeVars) {
@@ -869,7 +861,6 @@ namespace {
       }
 
       // Flatten the set of components.
-      SmallVector<Component, 1> flatComponents;
       flatComponents.reserve(
           representativeTypeVars.size() + cg.getOrphanedConstraints().size());
       for (auto rep: representativeTypeVars) {
@@ -934,10 +925,13 @@ namespace {
 
       // Reparent the type variable with the higher ID. The actual choice doesn't
       // matter, but this makes debugging easier.
-      if (rep1->getID() < rep2->getID())
+      if (rep1->getID() < rep2->getID()) {
+        validComponents.erase(rep2);
         representatives[rep2] = rep1;
-      else
+      } else {
+        validComponents.erase(rep1);
         representatives[rep1] = rep2;
+      }
       return true;
     }
 
@@ -948,6 +942,8 @@ namespace {
     TinyPtrVector<Constraint *> connectedComponents() {
       TinyPtrVector<Constraint *> oneWayConstraints;
 
+      auto &cs = cg.getConstraintSystem();
+
       // Perform a depth-first search from each type variable to identify
       // what component it is in.
       for (auto typeVar : typeVars) {
@@ -966,6 +962,10 @@ namespace {
               assert((inserted.second || inserted.first->second == typeVar) &&
                      "Wrong component?");
 
+              if (inserted.second)
+                if (!cs.getFixedType(found))
+                  validComponents.insert(typeVar);
+
               return inserted.second;
             },
             [&](Constraint *constraint) {
@@ -980,6 +980,17 @@ namespace {
             visitedConstraints);
       }
 
+      for (auto &constraint : cs.getConstraints()) {
+        if (constraint.getKind() == ConstraintKind::Disjunction ||
+            constraint.getKind() == ConstraintKind::Conjunction) {
+          for (auto typeVar : constraint.getTypeVariables()) {
+            auto rep = findRepresentative(typeVar);
+            if (validComponents.insert(rep).second)
+              ASSERT(cs.getFixedType(typeVar));
+          }
+        }
+      }
+
       return oneWayConstraints;
     }
 
@@ -1196,8 +1207,7 @@ namespace {
     /// solved after the components for type variables on the right-hand
     /// side of that constraint.
     SmallVector<TypeVariableType *, 4> computeOneWayComponentOrdering(
-        ArrayRef<TypeVariableType *> representativeTypeVars,
-        llvm::SmallDenseSet<TypeVariableType *> &validComponents) const {
+        ArrayRef<TypeVariableType *> representativeTypeVars) const {
       SmallVector<TypeVariableType *, 4> orderedReps;
       orderedReps.reserve(representativeTypeVars.size());
       SmallPtrSet<TypeVariableType *, 4> visited;