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

The added test case fails because the result builder inside `List2` is being type checked twice: Once for every overload of `List2`. Because of the way that result builders are being type checked right now, each overload of `List2` creates a new type variable for the key components in `foo` and the constraint system only keeps track of the key path component -> type mapping for the last solution discovered.

When we are now trying to look up the type of the first key path component for the first solution, the constraint system returns the type variable that is being used by the second solution and simplifying that type variable through the first solution fails because it doesn’t know about the type variable.

To fix the issue, make sure that the solutions keep track of the their type variables associated to key path components. That way the first solution owns its own key path component -> type variable mapping and is thus also able to simplify the type variable it has associated with the component.

Fixes rdar://80522345 [SR-14916]

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

@@ -8900,8 +8900,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 {
@@ -8913,6 +8914,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()
+}