Revert "Sema: Remove DependentComponentSplitterStep"

This reverts commit 9fb6d92.

Author: slavapestov

include/swift/Sema/ConstraintGraph.h

@@ -343,6 +343,12 @@ class ConstraintGraph {
     /// The constraints in this component.
     TinyPtrVector<Constraint *> constraints;
 
+    /// The set of components that this component depends on, such that
+    /// the partial solutions of the those components need to be available
+    /// before this component can be solved.
+    ///
+    SmallVector<unsigned, 2> dependencies;
+
   public:
     Component(unsigned solutionIndex) : solutionIndex(solutionIndex) { }
 
@@ -358,6 +364,11 @@ class ConstraintGraph {
       return constraints;
     }
 
+    /// Records a component which this component depends on.
+    void recordDependency(const Component &component);
+
+    ArrayRef<unsigned> getDependencies() const { return dependencies; }
+
     unsigned getNumDisjunctions() const { return numDisjunctions; }
   };
 

lib/Sema/CSStep.cpp

@@ -126,6 +126,7 @@ void SplitterStep::computeFollowupSteps(
   // Take the orphaned constraints, because they'll go into a component now.
   OrphanedConstraints = CG.takeOrphanedConstraints();
 
+  IncludeInMergedResults.resize(numComponents, true);
   Components.resize(numComponents);
   PartialSolutions = std::unique_ptr<SmallVector<Solution, 4>[]>(
       new SmallVector<Solution, 4>[numComponents]);
@@ -134,9 +135,26 @@ void SplitterStep::computeFollowupSteps(
   for (unsigned i : indices(components)) {
     unsigned solutionIndex = components[i].solutionIndex;
 
-    steps.push_back(std::make_unique<ComponentStep>(
-        CS, solutionIndex, &Components[i], std::move(components[i]),
-        PartialSolutions[solutionIndex]));
+    // If there are no dependencies, build a normal component step.
+    if (components[i].getDependencies().empty()) {
+      steps.push_back(std::make_unique<ComponentStep>(
+          CS, solutionIndex, &Components[i], std::move(components[i]),
+          PartialSolutions[solutionIndex]));
+      continue;
+    }
+
+    // Note that the partial results from any dependencies of this component
+    // need not be included in the final merged results, because they'll
+    // already be part of the partial results for this component.
+    for (auto dependsOn : components[i].getDependencies()) {
+      IncludeInMergedResults[dependsOn] = false;
+    }
+
+    // Otherwise, build a dependent component "splitter" step, which
+    // handles all combinations of incoming partial solutions.
+    steps.push_back(std::make_unique<DependentComponentSplitterStep>(
+        CS, &Components[i], solutionIndex, std::move(components[i]),
+        llvm::MutableArrayRef(PartialSolutions.get(), numComponents)));
   }
 
   assert(CS.InactiveConstraints.empty() && "Missed a constraint");
@@ -205,7 +223,8 @@ bool SplitterStep::mergePartialSolutions() const {
   SmallVector<unsigned, 2> countsVec;
   countsVec.reserve(numComponents);
   for (unsigned idx : range(numComponents)) {
-    countsVec.push_back(PartialSolutions[idx].size());
+    countsVec.push_back(
+        IncludeInMergedResults[idx] ? PartialSolutions[idx].size() : 1);
   }
 
   // Produce all combinations of partial solutions.
@@ -218,6 +237,9 @@ bool SplitterStep::mergePartialSolutions() const {
     // solutions.
     ConstraintSystem::SolverScope scope(CS);
     for (unsigned i : range(numComponents)) {
+      if (!IncludeInMergedResults[i])
+        continue;
+
       CS.replaySolution(PartialSolutions[i][indices[i]]);
     }
 
@@ -249,15 +271,77 @@ bool SplitterStep::mergePartialSolutions() const {
   return anySolutions;
 }
 
+StepResult DependentComponentSplitterStep::take(bool prevFailed) {
+  // "split" is considered a failure if previous step failed,
+  // or there is a failure recorded by constraint system, or
+  // system can't be simplified.
+  if (prevFailed || CS.getFailedConstraint() || CS.simplify())
+    return done(/*isSuccess=*/false);
+
+  // Figure out the sets of partial solutions that this component depends on.
+  SmallVector<const SmallVector<Solution, 4> *, 2> dependsOnSets;
+  for (auto index : Component.getDependencies()) {
+    dependsOnSets.push_back(&AllPartialSolutions[index]);
+  }
+
+  // Produce all combinations of partial solutions for the inputs.
+  SmallVector<std::unique_ptr<SolverStep>, 4> followup;
+  SmallVector<unsigned, 2> indices(Component.getDependencies().size(), 0);
+  auto dependsOnSetsRef = llvm::ArrayRef(dependsOnSets);
+  do {
+    // Form the set of input partial solutions.
+    SmallVector<const Solution *, 2> dependsOnSolutions;
+    for (auto index : swift::indices(indices)) {
+      dependsOnSolutions.push_back(&(*dependsOnSets[index])[indices[index]]);
+    }
+    ContextualSolutions.push_back(std::make_unique<SmallVector<Solution, 2>>());
+
+    followup.push_back(std::make_unique<ComponentStep>(
+        CS, Index, Constraints, Component, std::move(dependsOnSolutions),
+        *ContextualSolutions.back()));
+  } while (nextCombination(dependsOnSetsRef, indices));
+
+  /// Wait until all of the component steps are done.
+  return suspend(followup);
+}
+
+StepResult DependentComponentSplitterStep::resume(bool prevFailed) {
+  for (auto &ComponentStepSolutions : ContextualSolutions) {
+    Solutions.append(std::make_move_iterator(ComponentStepSolutions->begin()),
+                     std::make_move_iterator(ComponentStepSolutions->end()));
+  }
+  return done(/*isSuccess=*/!Solutions.empty());
+}
+
+void DependentComponentSplitterStep::print(llvm::raw_ostream &Out) {
+  Out << "DependentComponentSplitterStep for dependencies on [";
+  interleave(
+      Component.getDependencies(), [&](unsigned index) { Out << index; },
+      [&] { Out << ", "; });
+  Out << "]\n";
+}
+
 StepResult ComponentStep::take(bool prevFailed) {
   // One of the previous components created by "split"
   // failed, it means that we can't solve this component.
-  if (prevFailed || CS.isTooComplex(Solutions) || CS.worseThanBestSolution())
+  if ((prevFailed && DependsOnPartialSolutions.empty()) ||
+      CS.isTooComplex(Solutions) || CS.worseThanBestSolution())
     return done(/*isSuccess=*/false);
 
   // Setup active scope, only if previous component didn't fail.
   setupScope();
 
+  // If there are any dependent partial solutions to compose, do so now.
+  if (!DependsOnPartialSolutions.empty()) {
+    for (auto partial : DependsOnPartialSolutions) {
+      CS.replaySolution(*partial);
+    }
+
+    // Simplify again.
+    if (CS.failedConstraint || CS.simplify())
+      return done(/*isSuccess=*/false);
+  }
+
   /// Try to figure out what this step is going to be,
   /// after the scope has been established.
   SmallString<64> potentialBindings;

lib/Sema/CSStep.h

@@ -240,6 +240,10 @@ class SplitterStep final : public SolverStep {
 
   SmallVector<Constraint *, 4> OrphanedConstraints;
 
+  /// Whether to include the partial results of this component in the final
+  /// merged results.
+  SmallVector<bool, 4> IncludeInMergedResults;
+
 public:
   SplitterStep(ConstraintSystem &cs, SmallVectorImpl<Solution> &solutions)
       : SolverStep(cs, solutions) {}
@@ -265,6 +269,56 @@ class SplitterStep final : public SolverStep {
   bool mergePartialSolutions() const;
 };
 
+/// `DependentComponentSplitterStep` is responsible for composing the partial
+/// solutions from other components (on which this component depends) into
+/// the inputs based on which we can solve a particular component.
+class DependentComponentSplitterStep final : public SolverStep {
+  /// Constraints "in scope" of this step.
+  ConstraintList *Constraints;
+
+  /// Index into the parent splitter step.
+  unsigned Index;
+
+  /// The component that has dependencies.
+  ConstraintGraph::Component Component;
+
+  /// Array containing all of the partial solutions for the parent split.
+  MutableArrayRef<SmallVector<Solution, 4>> AllPartialSolutions;
+
+  /// The solutions computed the \c ComponentSteps created for each partial
+  /// solution combinations. Will be merged into the final \c Solutions vector
+  /// in \c resume.
+  std::vector<std::unique_ptr<SmallVector<Solution, 2>>> ContextualSolutions;
+
+  /// Take all of the constraints in this component and put them into
+  /// \c Constraints.
+  void injectConstraints() {
+    for (auto constraint : Component.getConstraints()) {
+      Constraints->erase(constraint);
+      Constraints->push_back(constraint);
+    }
+  }
+
+public:
+  DependentComponentSplitterStep(
+      ConstraintSystem &cs,
+      ConstraintList *constraints,
+      unsigned index,
+      ConstraintGraph::Component &&component,
+      MutableArrayRef<SmallVector<Solution, 4>> allPartialSolutions)
+    : SolverStep(cs, allPartialSolutions[index]), Constraints(constraints),
+      Index(index), Component(std::move(component)),
+      AllPartialSolutions(allPartialSolutions) {
+    assert(!Component.getDependencies().empty() && "Should use ComponentStep");
+    injectConstraints();
+  }
+
+  StepResult take(bool prevFailed) override;
+  StepResult resume(bool prevFailed) override;
+
+  void print(llvm::raw_ostream &Out) override;
+};
+
 
 /// `ComponentStep` represents a set of type variables and related
 /// constraints which could be solved independently. It's further
@@ -327,6 +381,10 @@ class ComponentStep final : public SolverStep {
   /// Constraints "in scope" of this step.
   ConstraintList *Constraints;
 
+  /// The set of partial solutions that should be composed before evaluating
+  /// this component.
+  SmallVector<const Solution *, 2> DependsOnPartialSolutions;
+
   /// Constraint which doesn't have any free type variables associated
   /// with it, which makes it disconnected in the graph.
   Constraint *OrphanedConstraint = nullptr;
@@ -361,6 +419,8 @@ class ComponentStep final : public SolverStep {
       constraints->erase(constraint);
       Constraints->push_back(constraint);
     }
+
+    assert(component.getDependencies().empty());
   }
 
   /// Create a component step that composes existing partial solutions before
@@ -369,11 +429,15 @@ class ComponentStep final : public SolverStep {
       ConstraintSystem &cs, unsigned index,
       ConstraintList *constraints,
       const ConstraintGraph::Component &component,
+      llvm::SmallVectorImpl<const Solution *> &&dependsOnPartialSolutions,
       SmallVectorImpl<Solution> &solutions)
         : SolverStep(cs, solutions), Index(index), IsSingle(false),
           OriginalScore(getCurrentScore()), OriginalBestScore(getBestScore()),
-          Constraints(constraints) {
+          Constraints(constraints),
+          DependsOnPartialSolutions(std::move(dependsOnPartialSolutions)) {
     TypeVars = component.typeVars;
+    assert(DependsOnPartialSolutions.size() ==
+           component.getDependencies().size());
 
     for (auto constraint : component.getConstraints()) {
       constraints->erase(constraint);

lib/Sema/ConstraintGraph.cpp

@@ -872,6 +872,10 @@ void ConstraintGraph::Component::addConstraint(Constraint *constraint) {
   constraints.push_back(constraint);
 }
 
+void ConstraintGraph::Component::recordDependency(const Component &component) {
+  dependencies.push_back(component.solutionIndex);
+}
+
 SmallVector<ConstraintGraph::Component, 1>
 ConstraintGraph::computeConnectedComponents(
            ArrayRef<TypeVariableType *> typeVars) {
@@ -1125,6 +1129,20 @@ void ConstraintGraph::printConnectedComponents(
                [&] {
                  out << ' ';
                });
+
+    auto dependencies = component.getDependencies();
+    if (dependencies.empty())
+      continue;
+
+    SmallVector<unsigned, 4> indices{dependencies.begin(), dependencies.end()};
+    // Sort dependencies so output is stable.
+    llvm::sort(indices);
+
+    // Print all of the one-way components.
+    out << " depends on ";
+    llvm::interleave(
+        indices, [&out](unsigned index) { out << index; },
+        [&out] { out << ", "; });
   }
 }