Sema: Split up gatherConstraints() into gatherAllConstraints() and gatherNearbyConstraints()

The two GatherKinds no longer share any implementation, so there's
no point keeping the logic together. Doing this also allows removing
the acceptConstraintFn from gatherAllConstraints(), which further
simplifies depthFirstSearch().

Author: slavapestov

include/swift/Sema/ConstraintGraph.h

@@ -295,25 +295,18 @@ class ConstraintGraph {
   /// to a type variable.
   void introduceToInference(TypeVariableType *typeVar, Type fixedType);
 
-  /// Describes which constraints \c gatherConstraints should gather.
-  enum class GatheringKind {
-    /// Gather constraints associated with all of the variables within the
-    /// same equivalence class as the given type variable, as well as its
-    /// immediate fixed bindings.
-    EquivalenceClass,
-    /// Gather all constraints that mention this type variable or type variables
-    /// that it is a fixed binding of. Unlike EquivalenceClass, this looks
-    /// through transitive fixed bindings. This can be used to find all the
-    /// constraints that may be affected when binding a type variable.
-    AllMentions,
-  };
+  /// Gather constraints associated with all of the variables within the
+  /// same equivalence class as the given type variable, as well as its
+  /// immediate fixed bindings.
+  llvm::TinyPtrVector<Constraint *>
+  gatherAllConstraints(TypeVariableType *typeVar);
 
-  /// Gather the set of constraints that involve the given type variable,
-  /// i.e., those constraints that will be affected when the type variable
-  /// gets merged or bound to a fixed type.
+  /// Gather all constraints that mention this type variable or type variables
+  /// that it is a fixed binding of. Unlike EquivalenceClass, this looks
+  /// through transitive fixed bindings. This can be used to find all the
+  /// constraints that may be affected when binding a type variable.
   llvm::TinyPtrVector<Constraint *>
-  gatherConstraints(TypeVariableType *typeVar,
-                    GatheringKind kind,
+  gatherNearbyConstraints(TypeVariableType *typeVar,
                     llvm::function_ref<bool(Constraint *)> acceptConstraint =
                         [](Constraint *constraint) { return true; });
 

lib/Sema/CSGen.cpp

@@ -400,8 +400,8 @@ namespace {
 
     llvm::SmallSetVector<ProtocolDecl *, 2> literalProtos;
     if (auto argTypeVar = argTy->getAs<TypeVariableType>()) {
-      auto constraints = CS.getConstraintGraph().gatherConstraints(
-          argTypeVar, ConstraintGraph::GatheringKind::EquivalenceClass,
+      auto constraints = CS.getConstraintGraph().gatherNearbyConstraints(
+          argTypeVar,
           [](Constraint *constraint) {
             return constraint->getKind() == ConstraintKind::LiteralConformsTo;
           });

lib/Sema/CSSimplify.cpp

@@ -15953,8 +15953,8 @@ ConstraintSystem::addKeyPathApplicationRootConstraint(Type root, ConstraintLocat
   if (!typeVar)
     return;
 
-  auto constraints = CG.gatherConstraints(
-      typeVar, ConstraintGraph::GatheringKind::EquivalenceClass,
+  auto constraints = CG.gatherNearbyConstraints(
+      typeVar,
       [&keyPathExpr](Constraint *constraint) -> bool {
         if (constraint->getKind() != ConstraintKind::KeyPath)
           return false;

lib/Sema/CSSolver.cpp

@@ -1876,8 +1876,8 @@ static Constraint *selectBestBindingDisjunction(
     if (!firstBindDisjunction)
       firstBindDisjunction = disjunction;
 
-    auto constraints = cs.getConstraintGraph().gatherConstraints(
-        typeVar, ConstraintGraph::GatheringKind::EquivalenceClass,
+    auto constraints = cs.getConstraintGraph().gatherNearbyConstraints(
+        typeVar,
         [](Constraint *constraint) {
           return constraint->getKind() == ConstraintKind::Conversion;
         });
@@ -1906,8 +1906,8 @@ ConstraintSystem::findConstraintThroughOptionals(
   while (visitedVars.insert(rep).second) {
     // Look for a disjunction that binds this type variable to an overload set.
     TypeVariableType *optionalObjectTypeVar = nullptr;
-    auto constraints = getConstraintGraph().gatherConstraints(
-        rep, ConstraintGraph::GatheringKind::EquivalenceClass,
+    auto constraints = getConstraintGraph().gatherNearbyConstraints(
+        rep,
         [&](Constraint *match) {
           // If we have an "optional object of" constraint, we may need to
           // look through it to find the constraint we're looking for.
@@ -2549,9 +2549,8 @@ void DisjunctionChoice::propagateConversionInfo(ConstraintSystem &cs) const {
     }
   }
 
-  auto constraints = cs.CG.gatherConstraints(
+  auto constraints = cs.CG.gatherNearbyConstraints(
       typeVar,
-      ConstraintGraph::GatheringKind::EquivalenceClass,
       [](Constraint *constraint) -> bool {
         switch (constraint->getKind()) {
         case ConstraintKind::Conversion:

lib/Sema/ConstraintGraph.cpp

@@ -545,18 +545,11 @@ void ConstraintGraph::retractBindings(TypeVariableType *typeVar,
 
 #pragma mark Algorithms
 
-/// Perform a depth-first search.
-///
-/// \param cg The constraint graph.
-/// \param typeVar The type variable we're searching from.
-/// \param visitConstraint Called before considering a constraint.
-/// \param visitedConstraints Set of already-visited constraints, used
-/// internally to avoid duplicated work.
 static void depthFirstSearch(
     ConstraintGraph &cg,
     TypeVariableType *typeVar,
-    llvm::function_ref<void(Constraint *)> visitConstraint,
     llvm::SmallPtrSet<TypeVariableType *, 4> &typeVars,
+    llvm::TinyPtrVector<Constraint *> &constraints,
     llvm::SmallPtrSet<Constraint *, 8> &visitedConstraints) {
   // If we're not looking at this type variable right now because we're
   // solving a conjunction element, don't consider its adjacencies.
@@ -570,11 +563,8 @@ static void depthFirstSearch(
   // Local function to visit adjacent type variables.
   auto visitAdjacencies = [&](ArrayRef<TypeVariableType *> adjTypeVars) {
     for (auto adj : adjTypeVars) {
-      if (adj == typeVar)
-        continue;
-
-      // Recurse into this node.
-      depthFirstSearch(cg, adj, visitConstraint, typeVars, visitedConstraints);
+      if (adj != typeVar)
+        depthFirstSearch(cg, adj, typeVars, constraints, visitedConstraints);
     }
   };
 
@@ -585,7 +575,7 @@ static void depthFirstSearch(
     if (!visitedConstraints.insert(constraint).second)
       continue;
 
-    visitConstraint(constraint);
+    constraints.push_back(constraint);
   }
 
   // Visit all of the other nodes in the equivalence class.
@@ -604,28 +594,22 @@ static void depthFirstSearch(
   visitAdjacencies(node.getReferencedVars());
 }
 
-llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherConstraints(
-    TypeVariableType *typeVar, GatheringKind kind,
-    llvm::function_ref<bool(Constraint *)> acceptConstraintFn) {
+llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherAllConstraints(
+    TypeVariableType *typeVar) {
   llvm::TinyPtrVector<Constraint *> constraints;
   llvm::SmallPtrSet<TypeVariableType *, 4> typeVars;
   llvm::SmallPtrSet<Constraint *, 8> visitedConstraints;
 
-  if (kind == GatheringKind::AllMentions) {
-    // If we've been asked for "all mentions" of a type variable, search for
-    // constraints involving both it and its fixed bindings.
-    depthFirstSearch(
-        *this, typeVar,
-        [&](Constraint *constraint) {
-          if (acceptConstraintFn(constraint))
-            constraints.push_back(constraint);
-        },
-        typeVars, visitedConstraints);
-    return constraints;
-  }
+  depthFirstSearch(*this, typeVar, typeVars, constraints, visitedConstraints);
+  return constraints;
+}
 
-  // Otherwise only search in the type var's equivalence class and immediate
-  // fixed bindings.
+llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherNearbyConstraints(
+    TypeVariableType *typeVar, 
+    llvm::function_ref<bool(Constraint *)> acceptConstraintFn) {
+  llvm::TinyPtrVector<Constraint *> constraints;
+  llvm::SmallPtrSet<TypeVariableType *, 4> typeVars;
+  llvm::SmallPtrSet<Constraint *, 8> visitedConstraints;
 
   // Local function to add constraints.
   auto addTypeVarConstraints = [&](TypeVariableType *adjTypeVar) {

lib/Sema/ConstraintSystem.cpp

@@ -267,8 +267,7 @@ void ConstraintSystem::assignFixedType(TypeVariableType *typeVar, Type type,
 void ConstraintSystem::addTypeVariableConstraintsToWorkList(
        TypeVariableType *typeVar) {
   // Activate the constraints affected by a change to this type variable.
-  auto gatheringKind = ConstraintGraph::GatheringKind::AllMentions;
-  for (auto *constraint : CG.gatherConstraints(typeVar, gatheringKind))
+  for (auto *constraint : CG.gatherAllConstraints(typeVar))
     if (!constraint->isActive())
       activateConstraint(constraint);
 }

lib/Sema/TypeOfReference.cpp

@@ -2058,8 +2058,8 @@ void ConstraintSystem::bindOverloadType(
       increaseScore(SK_KeyPathSubscript, locator);
 
       auto boundTypeVar = boundType->castTo<TypeVariableType>();
-      auto constraints = getConstraintGraph().gatherConstraints(
-          boundTypeVar, ConstraintGraph::GatheringKind::EquivalenceClass,
+      auto constraints = getConstraintGraph().gatherNearbyConstraints(
+          boundTypeVar,
           [](Constraint *constraint) {
             return constraint->getKind() == ConstraintKind::ApplicableFunction;
           });