From 65d8d29d2a1b4778097101c3aed1b1df963d3441 Mon Sep 17 00:00:00 2001
From: Rintaro Ishizaki <rishizaki@apple.com>
Date: Mon, 25 Mar 2019 09:54:31 -0700
Subject: [PATCH] [CodeCompletion] Infer keypath root type from the context
 type

So that this works:

```
func foo(_: KeyPath<SomeObject, Int>) {}
foo(\.<HERE>)
```

rdar://problem/46102807
---
 lib/IDE/CodeCompletion.cpp             | 27 ++++++++++++++++++--------
 test/IDE/complete_swift_key_path.swift | 26 ++++++++++++++++++-------
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/lib/IDE/CodeCompletion.cpp b/lib/IDE/CodeCompletion.cpp
index 313f244e573bd..5bddcd45f3261 100644
--- a/lib/IDE/CodeCompletion.cpp
+++ b/lib/IDE/CodeCompletion.cpp
@@ -5004,19 +5004,30 @@ void CodeCompletionCallbacksImpl::doneParsing() {
     bool OnRoot = !KPE->getComponents().front().isValid();
     Lookup.setIsSwiftKeyPathExpr(OnRoot);
 
-    auto ParsedType = BGT->getGenericArgs()[1];
-    auto Components = KPE->getComponents();
-    if (Components.back().getKind() ==
-        KeyPathExpr::Component::Kind::OptionalWrap) {
+    Type baseType = BGT->getGenericArgs()[OnRoot ? 0 : 1];
+    if (OnRoot && baseType->is<UnresolvedType>()) {
+      // Infer the root type of the keypath from the context type.
+      ExprContextInfo ContextInfo(CurDeclContext, ParsedExpr);
+      for (auto T : ContextInfo.getPossibleTypes()) {
+        if (auto unwrapped = T->getOptionalObjectType())
+          T = unwrapped;
+        if (!T->getAnyNominal() || !T->getAnyNominal()->getKeyPathTypeKind() ||
+            T->hasUnresolvedType() || !T->is<BoundGenericType>())
+          continue;
+        // Use the first KeyPath context type found.
+        baseType = T->castTo<BoundGenericType>()->getGenericArgs()[0];
+        break;
+      }
+    }
+    if (!OnRoot && KPE->getComponents().back().getKind() ==
+                       KeyPathExpr::Component::Kind::OptionalWrap) {
       // KeyPath expr with '?' (e.g. '\Ty.[0].prop?.another').
       // Althogh expected type is optional, we should unwrap it because it's
       // unwrapped.
-      ParsedType = ParsedType->getOptionalObjectType();
+      baseType = baseType->getOptionalObjectType();
     }
 
-    // The second generic type argument of KeyPath<Root, Value> should be
-    // the value we pull code completion results from.
-    Lookup.getValueExprCompletions(ParsedType);
+    Lookup.getValueExprCompletions(baseType);
     break;
   }
 
diff --git a/test/IDE/complete_swift_key_path.swift b/test/IDE/complete_swift_key_path.swift
index afe1963558aa4..f381cc4ec64e9 100644
--- a/test/IDE/complete_swift_key_path.swift
+++ b/test/IDE/complete_swift_key_path.swift
@@ -21,7 +21,9 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=EMPTY_2 | %FileCheck %s -check-prefix=INVALID
 
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_BASEONLY | %FileCheck %s -check-prefix=PERSONTYPE-DOT
-// FIXME: RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_EXPLICIT | %FileCheck %s -check-prefix=INVALID
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_EXPLICIT | %FileCheck %s -check-prefix=PERSONTYPE-DOT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_GENERIC_RESULT | %FileCheck %s -check-prefix=PERSONTYPE-DOT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_GENERIC_RESULT_OPTIONAL | %FileCheck %s -check-prefix=PERSONTYPE-DOT
 
 class Person {
     var name: String
@@ -131,9 +133,19 @@ let _ = \.#^EMPTY_1^#
 let _ = \.friends.#^EMPTY_2^#
 // INVALID-NOT: Begin completions
 
-let _: PartialKeyPath<Person> = \.#^CONTEXT_BASEONLY^#
-// Same as TYPE_DOT.
-
-let _: KeyPath<Person, String> = \.#^CONTEXT_EXPLICIT^#
-// FIXME: Currently, it's suggest nothing. Since we know the base type is
-// 'Person', we should suggest at least '.name'.
+func recvPartialKP(_ kp: PartialKeyPath<Person>) {
+  recvPartialKP(\.#^CONTEXT_BASEONLY^#)
+  // Same as TYPE_DOT.
+}
+func recvExplicitKP(_ kp: KeyPath<Person, String>) {
+  recvExplicitKP(\.#^CONTEXT_EXPLICIT^#)
+  // Same as TYPE_DOT.
+}
+func recvExplicitKPWithGenericResult<Result>(_ kp: KeyPath<Person, Result>) {
+  recvExplicitKPWithGenericResult(\.#^CONTEXT_GENERIC_RESULT^#)
+  // Same as TYPE_DOT.
+}
+func recvExplicitKPWithGenericResultOpt<Result>(_ kp: KeyPath<Person, Result>?) {
+  recvExplicitKPWithGenericResult(\.#^CONTEXT_GENERIC_RESULT_OPTIONAL^#
+  // Same as TYPE_DOT.
+}