Merge pull request #17791 from slavapestov/index-keypaths-4.2

IDE: Index references from keypaths [4.2]
diff --git a/lib/IDE/SourceEntityWalker.cpp b/lib/IDE/SourceEntityWalker.cpp
index 188cadf..e7deb88 100644
--- a/lib/IDE/SourceEntityWalker.cpp
+++ b/lib/IDE/SourceEntityWalker.cpp
@@ -62,6 +62,8 @@
   bool passModulePathElements(ArrayRef<ImportDecl::AccessPathElement> Path,
                               const clang::Module *ClangMod);
 
+  bool passReference(ValueDecl *D, Type Ty, SourceLoc Loc, SourceRange Range,
+                     ReferenceMetaData Data);
   bool passReference(ValueDecl *D, Type Ty, DeclNameLoc Loc, ReferenceMetaData Data);
   bool passReference(ModuleEntity Mod, std::pair<Identifier, SourceLoc> IdLoc);
 
@@ -324,6 +326,32 @@
       return { false, nullptr };
     return { false, E };
 
+  } else if (auto *KPE = dyn_cast<KeyPathExpr>(E)) {
+    for (auto &component : KPE->getComponents()) {
+      switch (component.getKind()) {
+      case KeyPathExpr::Component::Kind::Property:
+      case KeyPathExpr::Component::Kind::Subscript: {
+        auto *decl = component.getDeclRef().getDecl();
+        auto loc = component.getLoc();
+        SourceRange range(loc, loc);
+        passReference(decl, component.getComponentType(), loc, range,
+                      ReferenceMetaData(
+                        (isa<SubscriptDecl>(decl)
+                          ? SemaReferenceKind::SubscriptRef
+                          : SemaReferenceKind::DeclMemberRef),
+                        OpAccess));
+        break;
+      }
+
+      case KeyPathExpr::Component::Kind::Invalid:
+      case KeyPathExpr::Component::Kind::UnresolvedProperty:
+      case KeyPathExpr::Component::Kind::UnresolvedSubscript:
+      case KeyPathExpr::Component::Kind::OptionalChain:
+      case KeyPathExpr::Component::Kind::OptionalWrap:
+      case KeyPathExpr::Component::Kind::OptionalForce:
+        break;
+      }
+    }
   } else if (auto *BinE = dyn_cast<BinaryExpr>(E)) {
     // Visit in source order.
     if (!BinE->getArg()->getElement(0)->walk(*this))
@@ -459,31 +487,37 @@
 
 bool SemaAnnotator::
 passReference(ValueDecl *D, Type Ty, DeclNameLoc Loc, ReferenceMetaData Data) {
+  return passReference(D, Ty, Loc.getBaseNameLoc(), Loc.getSourceRange(), Data);
+}
+
+bool SemaAnnotator::
+passReference(ValueDecl *D, Type Ty, SourceLoc BaseNameLoc, SourceRange Range,
+              ReferenceMetaData Data) {
   TypeDecl *CtorTyRef = nullptr;
   ExtensionDecl *ExtDecl = nullptr;
 
   if (auto *TD = dyn_cast<TypeDecl>(D)) {
-    if (!CtorRefs.empty() && Loc.isValid()) {
+    if (!CtorRefs.empty() && BaseNameLoc.isValid()) {
       Expr *Fn = CtorRefs.back()->getFn();
-      if (Fn->getLoc() == Loc.getBaseNameLoc()) {
+      if (Fn->getLoc() == BaseNameLoc) {
         D = extractDecl(Fn);
         CtorTyRef = TD;
       }
     }
 
-    if (!ExtDecls.empty() && Loc.isValid()) {
+    if (!ExtDecls.empty() && BaseNameLoc.isValid()) {
       auto ExtTyLoc = ExtDecls.back()->getExtendedTypeLoc().getLoc();
-      if (ExtTyLoc.isValid() && ExtTyLoc == Loc.getBaseNameLoc()) {
+      if (ExtTyLoc.isValid() && ExtTyLoc == BaseNameLoc) {
         ExtDecl = ExtDecls.back();
       }
     }
   }
 
-  CharSourceRange Range =
+  CharSourceRange CharRange =
     Lexer::getCharSourceRangeFromSourceRange(D->getASTContext().SourceMgr,
-                                             Loc.getSourceRange());
-  bool Continue = SEWalker.visitDeclReference(D, Range, CtorTyRef, ExtDecl, Ty,
-                                              Data);
+                                             Range);
+  bool Continue = SEWalker.visitDeclReference(D, CharRange, CtorTyRef, ExtDecl,
+                                              Ty, Data);
   if (!Continue)
     Cancelled = true;
   return Continue;
diff --git a/test/Index/roles.swift b/test/Index/roles.swift
index adfe6b2..532c24a 100644
--- a/test/Index/roles.swift
+++ b/test/Index/roles.swift
@@ -483,3 +483,18 @@
   func foo() {}
   // CHECK: [[@LINE-1]]:8 | instance-method/Swift | {{.*}} | Def,RelChild | rel: 1
 }
+
+struct StructWithKeypath {
+  var x: Int = 0
+
+  subscript(idx: Int) -> Int { get { } set { } }
+}
+
+_ = \StructWithKeypath.x
+// CHECK: [[@LINE-1]]:24 | instance-property/Swift | x | s:14swift_ide_test17StructWithKeypathV1xSivp | Ref,Read | rel: 0
+// CHECK: [[@LINE-2]]:24 | instance-method/acc-get/Swift | getter:x | s:14swift_ide_test17StructWithKeypathV1xSivg | Ref,Call,Impl | rel: 0
+
+_ = \StructWithKeypath.[0]
+// CHECK: [[@LINE-1]]:24 | instance-property/subscript/Swift | subscript(_:) | s:14swift_ide_test17StructWithKeypathVyS2icip | Ref,Read | rel: 0
+// CHECK: [[@LINE-2]]:24 | instance-method/acc-get/Swift | getter:subscript(_:) | s:14swift_ide_test17StructWithKeypathVyS2icig | Ref,Call,Impl | rel: 0
+