[CodeCompletion] Update for SE-0249 KeyPath as functions

If the context type is a function type that receives one argument, infer
the key path expressions root type from the context.

rdar://problem/59302416

Author: rintaro

lib/IDE/CodeCompletion.cpp

@@ -5296,12 +5296,20 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       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 the context type is any of the KeyPath types, use it.
+        if (T->getAnyNominal() && T->getAnyNominal()->getKeyPathTypeKind() &&
+            !T->hasUnresolvedType() && T->is<BoundGenericType>()) {
+          baseType = T->castTo<BoundGenericType>()->getGenericArgs()[0];
+          break;
+        }
+
+        // KeyPath can be used as a function that receives its root type.
+        if (T->is<AnyFunctionType>() &&
+            T->castTo<AnyFunctionType>()->getNumParams() == 1) {
+          baseType = T->castTo<AnyFunctionType>()->getParams()[0].getOldType();
+          break;
+        }
       }
     }
     if (!OnRoot && KPE->getComponents().back().getKind() ==

test/IDE/complete_swift_key_path.swift

@@ -25,6 +25,11 @@
 // 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
 
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_FUNC | %FileCheck %s -check-prefix=PERSONTYPE-DOT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_FUNC_GENERICRESULT | %FileCheck %s -check-prefix=PERSONTYPE-DOT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_FUNC_ROOT | %FileCheck %s -check-prefix=PERSONTYPE-DOT
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=CONTEXT_FUNC_NONROOT | %FileCheck %s -check-prefix=OBJ-DOT
+
 class Person {
     var name: String
     var friends: [Person] = []
@@ -149,3 +154,19 @@ func recvExplicitKPWithGenericResultOpt<Result>(_ kp: KeyPath<Person, Result>?)
   recvExplicitKPWithGenericResult(\.#^CONTEXT_GENERIC_RESULT_OPTIONAL^#
   // Same as TYPE_DOT.
 }
+func recvFunc(_ fn: (Person) -> String) {
+  recvFunc(\.#^CONTEXT_FUNC^#)
+}
+func recvFuncGeneric<T>(_ fn: (Person) -> T) {
+  recvFunc(\.#^CONTEXT_FUNC_GENERICRESULT^#)
+}
+
+struct Wrap<T> {
+  func map<U>(_ fn: (T) -> U) -> U { fatalError() }
+}
+func testKeyPathAsFunctions(wrapped: Wrap<Person>) {
+  let _ = wrapped.map(\.#^CONTEXT_FUNC_ROOT^#)
+  // Same as TYPE_DOT.
+  let _ = wrapped.map(\.friends[0].#^CONTEXT_FUNC_NONROOT^#)
+  // Same as OBJ_DOT.
+}