[CodeCompletion] Provide basic code completion support for Swift KeyPath. rdar://31768743 (#9467)

Author: nkcsgexi

include/swift/IDE/CodeCompletion.h

@@ -503,6 +503,7 @@ enum class CompletionKind {
   ReturnStmtExpr,
   AfterPound,
   GenericParams,
+  SwiftKeyPath,
 };
 
 /// \brief A single code completion result.

include/swift/Parse/Parser.h

@@ -131,6 +131,8 @@ class Parser {
   bool InPoundLineEnvironment = false;
   bool InPoundIfEnvironment = false;
   bool InSwiftKeyPath = false;
+  Expr* SwiftKeyPathRoot = nullptr;
+  SourceLoc SwiftKeyPathSlashLoc = SourceLoc();
 
   LocalContext *CurLocalContext = nullptr;
 

lib/IDE/CodeCompletion.cpp

@@ -318,6 +318,16 @@ static bool KeyPathFilter(ValueDecl* decl, DeclVisibilityKind) {
          (isa<VarDecl>(decl) && decl->getDeclContext()->isTypeContext());
 }
 
+static bool SwiftKeyPathFilter(ValueDecl* decl, DeclVisibilityKind) {
+  switch(decl->getKind()){
+  case DeclKind::Var:
+  case DeclKind::Subscript:
+    return true;
+  default:
+    return false;
+  }
+}
+
 std::string swift::ide::removeCodeCompletionTokens(
     StringRef Input, StringRef TokenName, unsigned *CompletionOffset) {
   assert(TokenName.size() >= 1);
@@ -1601,6 +1611,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   bool IsSuperRefExpr = false;
   bool IsSelfRefExpr = false;
   bool IsKeyPathExpr = false;
+  bool IsSwiftKeyPathExpr = false;
   bool IsDynamicLookup = false;
   bool PreferFunctionReferencesToCalls = false;
   bool HaveLeadingSpace = false;
@@ -1766,6 +1777,10 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     IsKeyPathExpr = true;
   }
 
+  void setIsSwiftKeyPathExpr() {
+    IsSwiftKeyPathExpr = true;
+  }
+
   void setIsDynamicLookup() {
     IsDynamicLookup = true;
   }
@@ -2600,8 +2615,14 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     }
   }
 
+  bool shouldAddSubscriptCall() {
+    if (IsSwiftKeyPathExpr)
+      return true;
+    return !HaveDot;
+  }
+
   void addSubscriptCall(const SubscriptDecl *SD, DeclVisibilityKind Reason) {
-    assert(!HaveDot && "cannot add a subscript after a dot");
+    assert(shouldAddSubscriptCall() && "cannot add a subscript after a dot");
     CommandWordsPairs Pairs;
     CodeCompletionResultBuilder Builder(
         Sink,
@@ -2799,6 +2820,9 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
 
     if (IsKeyPathExpr && !KeyPathFilter(D, Reason))
       return;
+
+    if (IsSwiftKeyPathExpr && !SwiftKeyPathFilter(D, Reason))
+      return;
 
     if (!D->hasInterfaceType())
       TypeResolver->resolveDeclSignature(D);
@@ -2917,14 +2941,13 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
         addEnumElementRef(EED, Reason, /*HasTypeContext=*/false);
       }
 
-      if (HaveDot)
-        return;
-
-      if (auto *SD = dyn_cast<SubscriptDecl>(D)) {
-        if (ExprType->is<AnyMetatypeType>())
-          return;
-        addSubscriptCall(SD, Reason);
-        return;
+      // Swift key path allows .[0]
+      if (shouldAddSubscriptCall()) {
+        if (auto *SD = dyn_cast<SubscriptDecl>(D)) {
+          if (ExprType->is<AnyMetatypeType>())
+            return;
+          addSubscriptCall(SD, Reason);
+        }
       }
       return;
 
@@ -4393,6 +4416,8 @@ void CodeCompletionCallbacksImpl::completeDotExpr(Expr *E, SourceLoc DotLoc) {
     return;
 
   Kind = CompletionKind::DotExpr;
+  if (E->getKind() == ExprKind::KeyPath)
+    Kind = CompletionKind::SwiftKeyPath;
   if (ParseExprSelectorContext != ObjCSelectorContext::None) {
     PreferFunctionReferencesToCalls = true;
     CompleteExprSelectorContext = ParseExprSelectorContext;
@@ -4785,6 +4810,7 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
   case CompletionKind::GenericParams:
   case CompletionKind::KeyPathExpr:
   case CompletionKind::KeyPathExprDot:
+  case CompletionKind::SwiftKeyPath:
     break;
 
   case CompletionKind::StmtOrExpr:
@@ -5127,6 +5153,20 @@ void CodeCompletionCallbacksImpl::doneParsing() {
     break;
   }
 
+  case CompletionKind::SwiftKeyPath: {
+    Lookup.setHaveDot(DotLoc);
+    Lookup.setIsSwiftKeyPathExpr();
+    if (auto BGT = (*ExprType)->getAs<BoundGenericType>()) {
+      auto AllArgs = BGT->getGenericArgs();
+      if (AllArgs.size() == 2) {
+        // The second generic type argument of KeyPath<Root, Value> should be
+        // the value we pull code completion results from.
+        Lookup.getValueExprCompletions(AllArgs[1]);
+      }
+    }
+    break;
+  }
+
   case CompletionKind::StmtOrExpr:
     DoPostfixExprBeginning();
     break;

lib/Parse/ParseExpr.cpp

@@ -525,6 +525,7 @@ ParserResult<Expr> Parser::parseExprUnary(Diag<> Message, bool isExprBasic) {
 ParserResult<Expr> Parser::parseExprKeyPath() {
   // Consume '\'.
   SourceLoc backslashLoc = consumeToken(tok::backslash);
+  llvm::SaveAndRestore<SourceLoc> slashLoc(SwiftKeyPathSlashLoc, backslashLoc);
 
   // FIXME: diagnostics
   ParserResult<Expr> rootResult, pathResult;
@@ -538,6 +539,8 @@ ParserResult<Expr> Parser::parseExprKeyPath() {
   }
 
   if (startsWithSymbol(Tok, '.')) {
+    llvm::SaveAndRestore<Expr*> S(SwiftKeyPathRoot, rootResult.getPtrOrNull());
+
     // For uniformity, \.foo is parsed as if it were MAGIC.foo, so we need to
     // make sure the . is there, but parsing the ? in \.? as .? doesn't make
     // sense. This is all made more complicated by .?. being considered an
@@ -1136,8 +1139,18 @@ Parser::parseExprPostfixSuffix(ParserResult<Expr> Result, bool isExprBasic,
 
       // Handle "x.<tab>" for code completion.
       if (Tok.is(tok::code_complete)) {
-        if (CodeCompletion && Result.isNonNull())
+        if (CodeCompletion && Result.isNonNull()) {
+          if (InSwiftKeyPath) {
+            Result = makeParserResult(
+              new (Context) KeyPathExpr(SwiftKeyPathSlashLoc, Result.get(),
+                                        nullptr));
+          } else if (SwiftKeyPathRoot) {
+            Result = makeParserResult(
+              new (Context) KeyPathExpr(SwiftKeyPathSlashLoc, SwiftKeyPathRoot,
+                                        Result.get()));
+          }
           CodeCompletion->completeDotExpr(Result.get(), /*DotLoc=*/TokLoc);
+        }
         // Eat the code completion token because we handled it.
         consumeToken(tok::code_complete);
         Result.setHasCodeCompletion();

test/IDE/complete_swift_key_path.swift

@@ -0,0 +1,45 @@
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_1 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_2 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_3 | %FileCheck %s -check-prefix=PERSON-MEMBER
+
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_5 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_6 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_7 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_8 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_9 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_10 | %FileCheck %s -check-prefix=PERSON-MEMBER
+// RUN: %target-swift-ide-test -code-completion -source-filename %s -code-completion-token=PART_11 | %FileCheck %s -check-prefix=PERSON-MEMBER
+
+class Person {
+    var name: String
+    var friends: [Person] = []
+    var bestFriend: Person? = nil
+    init(name: String) {
+        self.name = name
+    }
+    func getName() -> String { return name }
+    subscript(_ index: Int) -> Int { get { return 1} }
+}
+
+let keyPath1 = \Person.#^PART_1^#
+let keyPath2 = \Person.friends[0].#^PART_2^#
+let keyPath3 = \Person.friends[0].friends[0].friends[0].#^PART_3^#
+
+// FIXME: the optionality keypath should work after our compiler is ready.
+let keyPath4 = \Person.bestFriend?.#^PART_4^#
+let keyPath5 = \Person.friends.[0].friends[0].friends[0].#^PART_5^#
+let keyPath6 = \[Person].[0].#^PART_6^#
+let keyPath7 = \[Person].[0].friends[0].#^PART_7^#
+
+func foo1(_ p : Person) {
+  _ = p[keyPath:\Person.#^PART_8^#]
+  _ = p[keyPath:\Person.friends[0].#^PART_9^#]
+  _ = p[keyPath:\[Person].[0].#^PART_10^#]
+  _ = p[keyPath:\Person.friends.[0].friends[0].friends[0].#^PART_11^#]
+}
+
+// PERSON-MEMBER: Begin completions, 4 items
+// PERSON-MEMBER-NEXT: Decl[InstanceVar]/CurrNominal:      name[#String#]; name=name
+// PERSON-MEMBER-NEXT: Decl[InstanceVar]/CurrNominal:      friends[#[Person]#]; name=friends
+// PERSON-MEMBER-NEXT: Decl[InstanceVar]/CurrNominal:      bestFriend[#Person?#]; name=bestFriend 
+// PERSON-MEMBER-NEXT: Decl[Subscript]/CurrNominal:        [{#Int#}][#Int#]; name=[Int]