[Constraint graph] Having "gathering all mentions" mean "gathering more mentions"

lib/Sema/CSBindings.cpp

@@ -394,7 +394,7 @@ ConstraintSystem::getPotentialBindings(TypeVariableType *typeVar) {
   // Gather the constraints associated with this type variable.
   auto constraints =
       getConstraintGraph().gatherConstraints(
-          typeVar, ConstraintGraph::GatheringKind::EquivalenceClass);
+          typeVar, ConstraintGraph::GatheringKind::PotentialBindings);
 
   PotentialBindings result(typeVar);
 

lib/Sema/CSSolver.cpp

@@ -1581,7 +1581,7 @@ void ConstraintSystem::ArgumentInfoCollector::walk(Type argType) {
     visited.insert(rep);
 
     auto constraints = CS.getConstraintGraph().gatherConstraints(
-        rep, ConstraintGraph::GatheringKind::EquivalenceClass);
+        rep, ConstraintGraph::GatheringKind::PotentialBindings);
 
     for (auto *constraint : constraints) {
       switch (constraint->getKind()) {

lib/Sema/ConstraintGraph.cpp

@@ -373,6 +373,90 @@ void ConstraintGraph::unbindTypeVariable(TypeVariableType *typeVar, Type fixed){
   }
 }
 
+#pragma mark Algorithms
+
+namespace {
+  /// When to visit the fixed bindings of a type variable.
+  enum class VisitFixedBindings {
+    /// Always visit the fixed bindings
+    Always,
+    /// Never visit the fixed bindings
+    Never,
+    /// Visit fixed bindings one level down, but no further.
+    Once,
+  };
+}
+
+/// Perform a depth-first search.
+///
+/// \param cg The constraint graph.
+/// \param typeVar The type variable we're searching from.
+/// \param preVisitNode Called before traversing a node. Must return \c
+/// false when the node has already been visited.
+/// \param visitConstraint Called before considering a constraint. If it
+/// returns \c false, that constraint will be skipped.
+/// \param visitedConstraints Set of already-visited constraints, used
+/// internally to avoid duplicated work.
+static void depthFirstSearch(
+    ConstraintGraph &cg,
+    TypeVariableType *typeVar,
+    llvm::function_ref<bool(TypeVariableType *)> preVisitNode,
+    llvm::function_ref<bool(Constraint *)> visitConstraint,
+    VisitFixedBindings visitFixedBindings,
+    llvm::SmallPtrSet<Constraint *, 8> &visitedConstraints) {
+  // Visit this node. If we've already seen it, bail out.
+  if (!preVisitNode(typeVar))
+    return;
+
+  // Local function to visit adjacent type variables.
+  auto visitAdjacencies = [&](ArrayRef<TypeVariableType *> adjTypeVars,
+                              VisitFixedBindings visitFixedBindings) {
+    for (auto adj : adjTypeVars) {
+      if (adj == typeVar)
+        continue;
+
+      // Recurse into this node.
+      depthFirstSearch(cg, adj, preVisitNode, visitConstraint,
+                       visitFixedBindings, visitedConstraints);
+    }
+  };
+
+  // Walk all of the constraints associated with this node to find related
+  // nodes.
+  auto &node = cg[typeVar];
+  for (auto constraint : node.getConstraints()) {
+    // If we've already seen this constraint, skip it.
+    if (!visitedConstraints.insert(constraint).second)
+      continue;
+
+    if (visitConstraint(constraint))
+      visitAdjacencies(constraint->getTypeVariables(), visitFixedBindings);
+  }
+
+  // Visit all of the other nodes in the equivalence class.
+  auto repTypeVar = cg.getConstraintSystem().getRepresentative(typeVar);
+  if (typeVar == repTypeVar) {
+    // We are the representative, so visit all of the other type variables
+    // in this equivalence class.
+    visitAdjacencies(node.getEquivalenceClass(), visitFixedBindings);
+  } else {
+    // We are not the representative; visit the representative.
+    visitAdjacencies(repTypeVar, visitFixedBindings);
+  }
+
+  // Walk any type variables related via fixed bindings.
+  switch (visitFixedBindings) {
+  case VisitFixedBindings::Always:
+    visitAdjacencies(node.getFixedBindings(), VisitFixedBindings::Always);
+      break;
+  case VisitFixedBindings::Never:
+      break;
+  case VisitFixedBindings::Once:
+    visitAdjacencies(node.getFixedBindings(), VisitFixedBindings::Never);
+      break;
+  }
+}
+
 llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherConstraints(
     TypeVariableType *typeVar, GatheringKind kind,
     llvm::function_ref<bool(Constraint *)> acceptConstraintFn) {
@@ -407,6 +491,42 @@ llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherConstraints(
         acceptConstraintFn(constraint);
   };
 
+  if (kind == GatheringKind::AllMentions ||
+      kind == GatheringKind::EquivalenceClass) {
+    // Perform a depth-first search in the constraint graph to find all of the
+    // constraints that could be affected.
+    VisitFixedBindings visitFixedBindings;
+    switch (kind) {
+    case GatheringKind::AllMentions:
+    case GatheringKind::PotentialBindings:
+      visitFixedBindings = VisitFixedBindings::Once;
+      break;
+
+    case GatheringKind::EquivalenceClass:
+      visitFixedBindings = VisitFixedBindings::Never;
+      break;
+    }
+
+    SmallPtrSet<TypeVariableType *, 4> typeVariables;
+    llvm::SmallPtrSet<Constraint *, 8> visitedConstraints;
+    depthFirstSearch(*this, typeVar,
+                     [&](TypeVariableType *typeVar) {
+                       return typeVariables.insert(typeVar).second;
+                     },
+                     [&](Constraint *constraint) {
+                       if (!shouldConsiderConstraint(constraint))
+                         return false;
+
+                       if (acceptConstraintFn(constraint))
+                         constraints.push_back(constraint);
+
+                       return kind == GatheringKind::AllMentions;
+                     },
+                     visitFixedBindings,
+                     visitedConstraints);
+    return constraints;
+  }
+
   // Add constraints for the given adjacent type variable.
   llvm::SmallPtrSet<TypeVariableType *, 4> typeVars;
   llvm::SmallPtrSet<Constraint *, 4> visitedConstraints;
@@ -435,90 +555,20 @@ llvm::TinyPtrVector<Constraint *> ConstraintGraph::gatherConstraints(
       if (visitedConstraints.insert(constraint).second &&
           acceptConstraint(constraint))
         constraints.push_back(constraint);
+    }
 
-      // If we want all mentions, visit type variables within each of our
+    if (kind == GatheringKind::PotentialBindings) {
+      // For any type variable mentioned in a fixed binding, add adjacent
       // constraints.
-      if (kind == GatheringKind::AllMentions) {
-        if (!shouldConsiderConstraint(constraint))
-          continue;
-
-        for (auto adjTypeVar : constraint->getTypeVariables()) {
-          addAdjacentConstraints(adjTypeVar);
-        }
+      for (auto adjTypeVar : node.getFixedBindings()) {
+        addAdjacentConstraints(adjTypeVar);
       }
     }
-
-    // For any type variable mentioned in a fixed binding, add adjacent
-    // constraints.
-    for (auto adjTypeVar : node.getFixedBindings()) {
-      addAdjacentConstraints(adjTypeVar);
-    }
   }
 
   return constraints;
 }
 
-#pragma mark Algorithms
-
-/// Perform a depth-first search.
-///
-/// \param cg The constraint graph.
-/// \param typeVar The type variable we're searching from.
-/// \param preVisitNode Called before traversing a node. Must return \c
-/// false when the node has already been visited.
-/// \param visitConstraint Called before considering a constraint. If it
-/// returns \c false, that constraint will be skipped.
-/// \param visitedConstraints Set of already-visited constraints, used
-/// internally to avoid duplicated work.
-static void depthFirstSearch(
-    ConstraintGraph &cg,
-    TypeVariableType *typeVar,
-    llvm::function_ref<bool(TypeVariableType *)> preVisitNode,
-    llvm::function_ref<bool(Constraint *)> visitConstraint,
-    llvm::SmallPtrSet<Constraint *, 8> &visitedConstraints) {
-  // Visit this node. If we've already seen it, bail out.
-  if (!preVisitNode(typeVar))
-    return;
-
-  // 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, preVisitNode, visitConstraint,
-                       visitedConstraints);
-    }
-  };
-
-  // Walk all of the constraints associated with this node to find related
-  // nodes.
-  auto &node = cg[typeVar];
-  for (auto constraint : node.getConstraints()) {
-    // If we've already seen this constraint, skip it.
-    if (!visitedConstraints.insert(constraint).second)
-      continue;
-
-    if (visitConstraint(constraint))
-      visitAdjacencies(constraint->getTypeVariables());
-  }
-
-  // Visit all of the other nodes in the equivalence class.
-  auto repTypeVar = cg.getConstraintSystem().getRepresentative(typeVar);
-  if (typeVar == repTypeVar) {
-    // We are the representative, so visit all of the other type variables
-    // in this equivalence class.
-    visitAdjacencies(node.getEquivalenceClass());
-  } else {
-    // We are not the representative; visit the representative.
-    visitAdjacencies(repTypeVar);
-  }
-
-  // Walk any type variables related via fixed bindings.
-  visitAdjacencies(node.getFixedBindings());
-}
-
 namespace {
   /// A union-find connected components algorithm used to find the connected
   /// components within a constraint graph.
@@ -787,6 +837,7 @@ namespace {
 
               return true;
             },
+            VisitFixedBindings::Always,
             visitedConstraints);
       }
 

lib/Sema/ConstraintGraph.h

@@ -181,6 +181,10 @@ class ConstraintGraph {
     /// Gather constraints associated with all of the variables within the
     /// same equivalence class as the given type variable.
     EquivalenceClass,
+    /// Gather constraints that might be useful for potential bindings,
+    /// which also involves gathering constraints for type variables found
+    /// via fixed bindings.
+    PotentialBindings,
     /// Gather all constraints that mention this type variable or type variables
     /// that it is equivalent to.
     AllMentions,

test/Constraints/keypath.swift

@@ -51,3 +51,25 @@ func testFunc() {
   let f = \S.i
   let _: (S) -> Int = f // expected-error {{cannot convert value of type 'KeyPath<S, Int>' to specified type '(S) -> Int'}}
 }
+
+// rdar://problem/54322807
+struct X<T> {
+  init(foo: KeyPath<T, Bool>) { }
+  init(foo: KeyPath<T, Bool?>) { }
+}
+
+struct Wibble {
+  var boolProperty = false
+}
+
+struct Bar {
+  var optWibble: Wibble? = nil
+}
+
+class Foo {
+  var optBar: Bar? = nil
+}
+
+func testFoo<T: Foo>(_: T) {
+  let _: X<T> = .init(foo: \.optBar!.optWibble?.boolProperty)
+}

validation-test/Sema/type_checker_perf/fast/rdar25866240.swift.gyb

@@ -1,4 +1,4 @@
-// RUN: %scale-test --begin 2 --end 10 --step 1 --select NumConstraintScopes %s -Xfrontend=-solver-disable-shrink -Xfrontend=-disable-constraint-solver-performance-hacks -Xfrontend=-solver-enable-operator-designated-types
+// RUN: %scale-test --begin 2 --end 15 --step 1 --select NumConstraintScopes %s -Xfrontend=-solver-disable-shrink -Xfrontend=-disable-constraint-solver-performance-hacks -Xfrontend=-solver-enable-operator-designated-types
 // REQUIRES: OS=macosx
 // REQUIRES: asserts