@@ -552,11 +552,18 @@ void ConstraintGraph::retractBindings(TypeVariableType *typeVar,
552552
553553#pragma mark Algorithms
554554
555+ // / Perform a depth-first search.
556+ // /
557+ // / \param cg The constraint graph.
558+ // / \param typeVar The type variable we're searching from.
559+ // / \param visitConstraint Called before considering a constraint.
560+ // / \param visitedConstraints Set of already-visited constraints, used
561+ // / internally to avoid duplicated work.
555562static void depthFirstSearch (
556563 ConstraintGraph &cg,
557564 TypeVariableType *typeVar,
565+ llvm::function_ref<void (Constraint *)> visitConstraint,
558566 llvm::SmallPtrSet<TypeVariableType *, 4> &typeVars,
559- llvm::TinyPtrVector<Constraint *> &constraints,
560567 llvm::SmallPtrSet<Constraint *, 8> &visitedConstraints) {
561568 // If we're not looking at this type variable right now because we're
562569 // solving a conjunction element, don't consider its adjacencies.
@@ -570,8 +577,11 @@ static void depthFirstSearch(
570577 // Local function to visit adjacent type variables.
571578 auto visitAdjacencies = [&](ArrayRef<TypeVariableType *> adjTypeVars) {
572579 for (auto adj : adjTypeVars) {
573- if (adj != typeVar)
574- depthFirstSearch (cg, adj, typeVars, constraints, visitedConstraints);
580+ if (adj == typeVar)
581+ continue ;
582+
583+ // Recurse into this node.
584+ depthFirstSearch (cg, adj, visitConstraint, typeVars, visitedConstraints);
575585 }
576586 };
577587
@@ -582,7 +592,7 @@ static void depthFirstSearch(
582592 if (!visitedConstraints.insert (constraint).second )
583593 continue ;
584594
585- constraints. push_back (constraint);
595+ visitConstraint (constraint);
586596 }
587597
588598 // Visit all of the other nodes in the equivalence class.
@@ -601,22 +611,28 @@ static void depthFirstSearch(
601611 visitAdjacencies (node.getReferencedVars ());
602612}
603613
604- llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherAllConstraints (
605- TypeVariableType *typeVar) {
614+ llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherConstraints (
615+ TypeVariableType *typeVar, GatheringKind kind,
616+ llvm::function_ref<bool (Constraint *)> acceptConstraintFn) {
606617 llvm::TinyPtrVector<Constraint *> constraints;
607618 llvm::SmallPtrSet<TypeVariableType *, 4 > typeVars;
608619 llvm::SmallPtrSet<Constraint *, 8 > visitedConstraints;
609620
610- depthFirstSearch (*this , typeVar, typeVars, constraints, visitedConstraints);
611- return constraints;
612- }
621+ if (kind == GatheringKind::AllMentions) {
622+ // If we've been asked for "all mentions" of a type variable, search for
623+ // constraints involving both it and its fixed bindings.
624+ depthFirstSearch (
625+ *this , typeVar,
626+ [&](Constraint *constraint) {
627+ if (acceptConstraintFn (constraint))
628+ constraints.push_back (constraint);
629+ },
630+ typeVars, visitedConstraints);
631+ return constraints;
632+ }
613633
614- llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherNearbyConstraints (
615- TypeVariableType *typeVar,
616- llvm::function_ref<bool (Constraint *)> acceptConstraintFn) {
617- llvm::TinyPtrVector<Constraint *> constraints;
618- llvm::SmallPtrSet<TypeVariableType *, 4 > typeVars;
619- llvm::SmallPtrSet<Constraint *, 8 > visitedConstraints;
634+ // Otherwise only search in the type var's equivalence class and immediate
635+ // fixed bindings.
620636
621637 // Local function to add constraints.
622638 auto addTypeVarConstraints = [&](TypeVariableType *adjTypeVar) {
0 commit comments