[Constraint graph] Move component sorting into connected components.

(cherry picked from commit dec149c)

Author: DougGregor

lib/Sema/CSStep.cpp

@@ -125,19 +125,11 @@ void SplitterStep::computeFollowupSteps(
 
   // Add components.
   for (unsigned i : indices(components)) {
+    unsigned solutionIndex = components[i].solutionIndex;
     componentSteps.push_back(llvm::make_unique<ComponentStep>(
-        CS, i, &Components[i], std::move(components[i]), PartialSolutions[i]));
+        CS, solutionIndex, &Components[i], std::move(components[i]),
+        PartialSolutions[solutionIndex]));
   }
-
-  // Create component ordering based on the information associated
-  // with constraints in each step - e.g. number of disjunctions,
-  // since components are going to be executed in LIFO order, we'd
-  // want to have smaller/faster components at the back of the list.
-  std::sort(componentSteps.begin(), componentSteps.end(),
-            [](const std::unique_ptr<ComponentStep> &lhs,
-               const std::unique_ptr<ComponentStep> &rhs) {
-              return lhs->disjunctionCount() > rhs->disjunctionCount();
-            });
 }
 
 bool SplitterStep::mergePartialSolutions() const {

lib/Sema/CSStep.h

@@ -341,10 +341,6 @@ class ComponentStep final : public SolverStep {
   /// Constraints "in scope" of this step.
   ConstraintList *Constraints;
 
-  /// Number of disjunction constraints associated with this step,
-  /// used to aid in ordering of the components.
-  unsigned NumDisjunctions = 0;
-
   /// Constraint which doesn't have any free type variables associated
   /// with it, which makes it disconnected in the graph.
   Constraint *OrphanedConstraint = nullptr;
@@ -367,36 +363,23 @@ class ComponentStep final : public SolverStep {
         OriginalScore(getCurrentScore()), OriginalBestScore(getBestScore()),
         Constraints(constraints) {
     if (component.isOrphaned()) {
-      assert(component.constraints.size() == 1);
-      OrphanedConstraint = component.constraints.front();
+      assert(component.getConstraints().size() == 1);
+      OrphanedConstraint = component.getConstraints().front();
     } else {
       assert(component.typeVars.size() > 0);
     }
 
     TypeVars = std::move(component.typeVars);
 
-    for (auto constraint : component.constraints) {
+    for (auto constraint : component.getConstraints()) {
       constraints->erase(constraint);
-      record(constraint);
+      Constraints->push_back(constraint);
     }
   }
 
-private:
-  /// Record a constraint as associated with this step.
-  void record(Constraint *constraint) {
-    Constraints->push_back(constraint);
-    if (constraint->getKind() == ConstraintKind::Disjunction)
-      ++NumDisjunctions;
-  }
-
-public:
-
   StepResult take(bool prevFailed) override;
   StepResult resume(bool prevFailed) override;
 
-  // The number of disjunction constraints associated with this component.
-  unsigned disjunctionCount() const { return NumDisjunctions; }
-
   void print(llvm::raw_ostream &Out) override {
     Out << "ComponentStep with at #" << Index << '\n';
   }

lib/Sema/ConstraintGraph.cpp

@@ -556,7 +556,7 @@ namespace {
           // We haven't allocated this component yet; do so now.
           knownComponentIdx = componentIdxMap.insert(
               {rep, componentIdxMap.size()}).first;
-          components.push_back({ });
+          components.push_back(Component(components.size()));
         }
 
         // Record this type variabgetConstraintsle as part of the component.
@@ -576,13 +576,24 @@ namespace {
         auto typeVar = constraintTypeVars.front();
         auto rep = findRepresentative(typeVar);
         assert(componentIdxMap.count(rep) > 0);
-        components[componentIdxMap[rep]].constraints.push_back(&constraint);
+        components[componentIdxMap[rep]].addConstraint(&constraint);
       }
 
       // Gather orphaned constraints; each gets its own component.
       for (auto orphaned : cg.getOrphanedConstraints()) {
-        components.push_back({ });
-        components.back().constraints.push_back(orphaned);
+        components.push_back(Component(components.size()));
+        components.back().addConstraint(orphaned);
+      }
+
+      // Create component ordering based on the information associated
+      // with constraints in each step - e.g. number of disjunctions,
+      // since components are going to be executed in LIFO order, we'd
+      // want to have smaller/faster components at the back of the list.
+      if (components.size() > 1) {
+        std::sort(components.begin(), components.end(),
+                  [&](const Component &lhs, const Component &rhs) {
+                    return lhs.getNumDisjunctions() > rhs.getNumDisjunctions();
+                  });
       }
 
       return components;
@@ -637,6 +648,13 @@ namespace {
   };
 }
 
+void ConstraintGraph::Component::addConstraint(Constraint *constraint) {
+  if (constraint->getKind() == ConstraintKind::Disjunction)
+    ++numDisjunctions;
+
+  constraints.push_back(constraint);
+}
+
 SmallVector<ConstraintGraph::Component, 1>
 ConstraintGraph::computeConnectedComponents(
            ArrayRef<TypeVariableType *> typeVars) {

lib/Sema/ConstraintGraph.h

@@ -209,13 +209,37 @@ class ConstraintGraph {
     /// The type variables in this component.
     TinyPtrVector<TypeVariableType *> typeVars;
 
+    /// The original index of this component in the list of components,
+    /// used to provide the index of where the partial solutions will occur.
+    /// FIXME: This is needed due to some ordering dependencies in the
+    /// merging of partial solutions, which appears to also be related
+    /// DisjunctionStep::pruneOverloads() short-circuiting. It should be
+    /// removed.
+    unsigned solutionIndex;
+
+  private:
+    /// The number of disjunctions in this component.
+    unsigned numDisjunctions = 0;
+
     /// The constraints in this component.
     TinyPtrVector<Constraint *> constraints;
 
+  public:
+    Component(unsigned solutionIndex) : solutionIndex(solutionIndex) { }
+
     /// Whether this component represents an orphaned constraint.
     bool isOrphaned() const {
       return typeVars.empty();
     }
+
+    /// Add a constraint.
+    void addConstraint(Constraint *constraint);
+
+    const TinyPtrVector<Constraint *> &getConstraints() const {
+      return constraints;
+    }
+
+    unsigned getNumDisjunctions() const { return numDisjunctions; }
   };
 
   /// Compute the connected components of the graph.