[CSStep] Record all of the type vars related to component while creating a split

Before this change some of the type variables, especially the ones
that have been previously bound haven't been associated with any
of the component steps, so it had to be done by scope once per step,
but now such recording is done once per split.

Author: xedin

lib/Sema/CSStep.cpp

@@ -30,13 +30,8 @@ ComponentStep::Scope::Scope(ComponentStep &component)
     : CS(component.CS), Component(component) {
   TypeVars = std::move(CS.TypeVariables);
 
-  for (auto *typeVar : component.TypeVars) {
-    // Include type variables that either belong to this
-    // component or have been bound.
-    if (component.TypeVars.count(typeVar) > 0 ||
-        typeVar->getImpl().getFixedType(nullptr))
-      CS.TypeVariables.push_back(typeVar);
-  }
+  for (auto *typeVar : component.TypeVars)
+    CS.TypeVariables.push_back(typeVar);
 
   auto &workList = CS.InactiveConstraints;
   workList.splice(workList.end(), *component.Constraints);
@@ -132,21 +127,22 @@ void SplitterStep::computeFollowupSteps(
   }
 
   // Map type variables and constraints into appropriate steps.
+  llvm::DenseMap<TypeVariableType *, unsigned> typeVarComponent;
   llvm::DenseMap<Constraint *, unsigned> constraintComponent;
   for (unsigned i = 0, n = typeVars.size(); i != n; ++i) {
     auto *typeVar = typeVars[i];
-    auto *component = componentSteps[components[i]];
+    // Record the component of this type variable.
+    typeVarComponent[typeVar] = components[i];
 
-    component->record(typeVar);
     for (auto *constraint : CG[typeVar].getConstraints())
       constraintComponent[constraint] = components[i];
   }
 
   // Add the orphaned components to the mapping from constraints to components.
-  unsigned firstOrphanedConstraint =
+  unsigned firstOrphanedComponent =
       numComponents - CG.getOrphanedConstraints().size();
   {
-    unsigned component = firstOrphanedConstraint;
+    unsigned component = firstOrphanedComponent;
     for (auto *constraint : CG.getOrphanedConstraints()) {
       // Register this orphan constraint both as associated with
       // a given component as a regular constrant, as well as an
@@ -157,6 +153,22 @@ void SplitterStep::computeFollowupSteps(
     }
   }
 
+  for (auto *typeVar : CS.TypeVariables) {
+    auto known = typeVarComponent.find(typeVar);
+    // If current type variable is associated with
+    // a certain component step, record it as being so.
+    if (known != typeVarComponent.end()) {
+      componentSteps[known->second]->record(typeVar);
+      continue;
+    }
+
+    // Otherwise, associate it with all of the component steps,
+    // expect for components with orphaned constraints, they are
+    // not supposed to have any type variables.
+    for (unsigned i = 0; i != firstOrphanedComponent; ++i)
+      componentSteps[i]->record(typeVar);
+  }
+
   // Transfer all of the constraints from the work list to
   // the appropriate component.
   auto &workList = CS.InactiveConstraints;

lib/Sema/CSStep.h

@@ -23,7 +23,6 @@
 #include "swift/AST/Types.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
-#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include <memory>
 
@@ -311,7 +310,7 @@ class ComponentStep final : public SolverStep {
   std::unique_ptr<Scope> ComponentScope = nullptr;
 
   /// Type variables and constraints "in scope" of this step.
-  SmallPtrSet<TypeVariableType *, 16> TypeVars;
+  SmallVector<TypeVariableType *, 16> TypeVars;
   /// Constraints "in scope" of this step.
   ConstraintList *Constraints;
 
@@ -332,7 +331,7 @@ class ComponentStep final : public SolverStep {
 
 public:
   /// Record a type variable as associated with this step.
-  void record(TypeVariableType *typeVar) { TypeVars.insert(typeVar); }
+  void record(TypeVariableType *typeVar) { TypeVars.push_back(typeVar); }
 
   /// Record a constraint as associated with this step.
   void record(Constraint *constraint) {
@@ -349,15 +348,14 @@ class ComponentStep final : public SolverStep {
   }
 
   void setup() override {
-    // If this component has oprhaned constraint attached,
-    // let's return it ot the graph.
-    if (OrphanedConstraint)
-      CS.CG.setOrphanedConstraint(OrphanedConstraint);
-
     // If this is a single component, there is
     // no need to preliminary modify constraint system.
-    if (!IsSingle)
+    if (!IsSingle) {
       ComponentScope = llvm::make_unique<Scope>(*this);
+      // If this component has oprhaned constraint attached,
+      // let's return it ot the graph.
+      CS.CG.setOrphanedConstraint(OrphanedConstraint);
+    }
   }
 
   StepResult take(bool prevFailed) override;