[CodeCompletion] Handle "KeyPath as function" thunk in SanitizeExpr

rdar://problem/60982638

Author: rintaro

lib/Sema/CSGen.cpp

@@ -914,6 +914,35 @@ namespace {
   };
 } // end anonymous namespace
 
+namespace {
+// Check if \p E is a call expression to curried thunk of "KeyPath as function".
+// i.e. '{ `$kp$` in { $0[keyPath: $kp$] } }(keypath)'
+static bool isKeyPathCurriedThunkCallExpr(Expr *E) {
+  auto CE = dyn_cast<CallExpr>(E);
+  if (!CE)
+    return false;
+  auto thunk = dyn_cast<AutoClosureExpr>(CE->getFn());
+  if (!thunk)
+    return false;
+  if (thunk->getParameters()->size() != 1 ||
+      thunk->getParameters()->get(0)->getParameterName().str() != "$kp$")
+    return false;
+
+  auto PE = dyn_cast<ParenExpr>(CE->getArg());
+  if (!PE)
+    return false;
+  return isa<KeyPathExpr>(PE->getSubExpr());
+}
+
+// Extract the keypath expression from the curried thunk expression.
+static Expr *extractKeyPathFromCurryThunkCall(Expr *E) {
+  assert(isKeyPathCurriedThunkCallExpr(E));
+  auto call = cast<CallExpr>(E);
+  auto arg = cast<ParenExpr>(call->getArg());
+  return arg->getSubExpr();
+}
+} // end anonymous namespace
+
 namespace {
 
   class ConstraintGenerator : public ExprVisitor<ConstraintGenerator, Type> {
@@ -3787,13 +3816,20 @@ namespace {
           continue;
         }
 
+        // Extract keypath from '{ `$kp$` in { $0[keyPath: $kp$] } }(keypath)'
+        if (isKeyPathCurriedThunkCallExpr(expr)) {
+          expr = extractKeyPathFromCurryThunkCall(expr);
+          continue;
+        }
+
         // Restore '@autoclosure'd value.
         if (auto ACE = dyn_cast<AutoClosureExpr>(expr)) {
           // This is only valid if the closure doesn't have parameters.
           if (ACE->getParameters()->size() == 0) {
             expr = ACE->getSingleExpressionBody();
             continue;
           }
+          llvm_unreachable("other AutoClosureExpr must be handled specially");
         }
 
         // Remove any semantic expression injected by typechecking.

test/IDE/complete_call_arg.swift

@@ -104,6 +104,8 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TUPLEELEM_1 | %FileCheck %s -check-prefix=TUPLEELEM_1
 // RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=TUPLEELEM_2 | %FileCheck %s -check-prefix=TUPLEELEM_2
 
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=KEYPATH_THUNK_BASE | %FileCheck %s -check-prefix=KEYPATH_THUNK_BASE
+
 var i1 = 1
 var i2 = 2
 var oi1 : Int?
@@ -829,3 +831,20 @@ func testTupleElement(arg: (SimpleEnum, SimpleEnum)) {
   testTupleElement(arg: (.foo, .bar, .#^TUPLEELEM_2^#))
 // TUPLEELEM_2-NOT: Begin completions
 }
+
+func testKeyPathThunkInBase() {
+    struct TestKP {
+        var value: Int { 1 }
+    }
+    struct TestKPResult {
+        func testFunc(_ arg: SimpleEnum) {}
+    }
+    func foo(_ fn: (TestKP) -> Int) -> TestKPResult { TestKPResult() }
+
+    foo(\.value).testFunc(.#^KEYPATH_THUNK_BASE^#)
+// KEYPATH_THUNK_BASE: Begin completions, 3 items
+// KEYPATH_THUNK_BASE-DAG: Decl[EnumElement]/ExprSpecific/TypeRelation[Identical]:     foo[#SimpleEnum#]; name=foo
+// KEYPATH_THUNK_BASE-DAG: Decl[EnumElement]/ExprSpecific/TypeRelation[Identical]:     bar[#SimpleEnum#]; name=bar
+// KEYPATH_THUNK_BASE-DAG: Decl[EnumElement]/ExprSpecific/TypeRelation[Identical]:     baz[#SimpleEnum#]; name=baz
+// KEYPATH_THUNK_BASE: End completions
+}