Merge pull request #38389 from ahoppen/pr/record-keypath-component-types-in-solution

[CodeCompletion] Record key path component types in the constraint system solution

Author: ahoppen

include/swift/Sema/ConstraintSystem.h

@@ -1185,6 +1185,10 @@ class Solution {
   /// The node -> type mappings introduced by this solution.
   llvm::DenseMap<ASTNode, Type> nodeTypes;
 
+  /// The key path component types introduced by this solution.
+  llvm::DenseMap<std::pair<const KeyPathExpr *, unsigned>, TypeBase *>
+      keyPathComponentTypes;
+
   /// Contextual types introduced by this solution.
   std::vector<std::pair<ASTNode, ContextualTypeInfo>> contextualTypes;
 
@@ -1300,6 +1304,9 @@ class Solution {
   /// Retrieve the type of the given node, as recorded in this solution.
   Type getType(ASTNode node) const;
 
+  /// Retrieve the type of the \p ComponentIndex-th component in \p KP.
+  Type getType(const KeyPathExpr *KP, unsigned ComponentIndex) const;
+
   /// Retrieve the type of the given node as recorded in this solution
   /// and resolve all of the type variables in contains to form a fully
   /// "resolved" concrete type.

lib/Sema/CSApply.cpp

@@ -8901,8 +8901,9 @@ bool Solution::hasType(ASTNode node) const {
 }
 
 bool Solution::hasType(const KeyPathExpr *KP, unsigned ComponentIndex) const {
-  auto &cs = getConstraintSystem();
-  return cs.hasType(KP, ComponentIndex);
+  assert(KP && "Expected non-null key path parameter!");
+  return keyPathComponentTypes.find(std::make_pair(KP, ComponentIndex))
+            != keyPathComponentTypes.end();
 }
 
 Type Solution::getType(ASTNode node) const {
@@ -8914,6 +8915,11 @@ Type Solution::getType(ASTNode node) const {
   return cs.getType(node);
 }
 
+Type Solution::getType(const KeyPathExpr *KP, unsigned I) const {
+  assert(hasType(KP, I) && "Expected type to have been set!");
+  return keyPathComponentTypes.find(std::make_pair(KP, I))->second;
+}
+
 Type Solution::getResolvedType(ASTNode node) const {
   return simplifyType(getType(node));
 }

lib/Sema/CSSolver.cpp

@@ -175,6 +175,9 @@ Solution ConstraintSystem::finalize() {
   for (auto &nodeType : NodeTypes) {
     solution.nodeTypes.insert(nodeType);
   }
+  for (auto &keyPathComponentType : KeyPathComponentTypes) {
+    solution.keyPathComponentTypes.insert(keyPathComponentType);
+  }
 
   // Remember contextual types.
   solution.contextualTypes.assign(

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -1347,7 +1347,7 @@ void KeyPathTypeCheckCompletionCallback::sawSolution(
   } else {
     // We are completing after a component. Get the previous component's result
     // type.
-    BaseType = S.simplifyType(CS.getType(KeyPath, ComponentIndex - 1));
+    BaseType = S.simplifyType(S.getType(KeyPath, ComponentIndex - 1));
   }
 
   // If ExpectedTy is a duplicate of any other result, ignore this solution.

test/IDE/complete_sr14916.swift

@@ -0,0 +1,29 @@
+// RUN: %swift-ide-test -code-completion -code-completion-token COMPLETE -source-filename %s | %FileCheck %s
+
+struct Foo {
+	var bar: Int
+}
+
+protocol View2 {}
+struct EmptyView: View2 {}
+
+@resultBuilder public struct ViewBuilder2 {
+  public static func buildBlock(_ content: EmptyView) -> EmptyView { fatalError() }
+}
+
+public struct List2 {
+    public init(selection: Int?, @ViewBuilder2 content: () -> EmptyView)
+    public init(selection: String?, @ViewBuilder2 content: () -> EmptyView)
+}
+
+func foo(kp: (Foo) -> String) {}
+
+func foo() {
+    List2 {
+        foo(kp: \.self#^COMPLETE^#)
+// CHECK:      Begin completions, 1 items
+// CHECK-NEXT: Decl[InstanceVar]/CurrNominal:      .bar[#Int#];
+// CHECK-NEXT: End completions
+    }
+    .unknownMethod()
+}