Merge pull request #23542 from rintaro/ide-completion-keypathcontext-rdar46102807

[CodeCompletion] Infer keypath root type from the context type

Author: rintaro

lib/IDE/CodeCompletion.cpp

@@ -4992,19 +4992,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;
   }
 

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.
+}