[clangd] Use canonical declarations in ReferenceFinder.

Summary:
handleDeclOccurrencce reports a canonical declartion, so stick to use
canonical declarations to determine whether a declaration is in the
target set.

Also fix a previous ref test which misses a matched label (it fails without this
patch).

Reviewers: sammccall

Subscribers: ilya-biryukov, ioeric, MaskRay, jkorous, arphaman, kadircet, cfe-commits

Differential Revision: https://reviews.llvm.org/D52871

2 files changed, 39 insertions, 29 deletions

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index da98d88f713..69c6aad21a1 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -361,7 +361,7 @@ namespace {
 class ReferenceFinder : public index::IndexDataConsumer {
 public:
   struct Reference {
-    const Decl *Target;
+    const Decl *CanonicalTarget;
     SourceLocation Loc;
     index::SymbolRoleSet Role;
   };
@@ -370,22 +370,23 @@ public:
                   const std::vector<const Decl *> &TargetDecls)
       : AST(AST) {
     for (const Decl *D : TargetDecls)
-      Targets.insert(D);
+      CanonicalTargets.insert(D->getCanonicalDecl());
   }
 
   std::vector<Reference> take() && {
     std::sort(References.begin(), References.end(),
               [](const Reference &L, const Reference &R) {
-                return std::tie(L.Loc, L.Target, L.Role) <
-                       std::tie(R.Loc, R.Target, R.Role);
+                return std::tie(L.Loc, L.CanonicalTarget, L.Role) <
+                       std::tie(R.Loc, R.CanonicalTarget, R.Role);
               });
     // We sometimes see duplicates when parts of the AST get traversed twice.
-    References.erase(std::unique(References.begin(), References.end(),
-                                 [](const Reference &L, const Reference &R) {
-                                   return std::tie(L.Target, L.Loc, L.Role) ==
-                                          std::tie(R.Target, R.Loc, R.Role);
-                                 }),
-                     References.end());
+    References.erase(
+        std::unique(References.begin(), References.end(),
+                    [](const Reference &L, const Reference &R) {
+                      return std::tie(L.CanonicalTarget, L.Loc, L.Role) ==
+                             std::tie(R.CanonicalTarget, R.Loc, R.Role);
+                    }),
+        References.end());
     return std::move(References);
   }
 
@@ -394,15 +395,16 @@ public:
                       ArrayRef<index::SymbolRelation> Relations,
                       SourceLocation Loc,
                       index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
+    assert(D->isCanonicalDecl() && "expect D to be a canonical declaration");
     const SourceManager &SM = AST.getSourceManager();
     Loc = SM.getFileLoc(Loc);
-    if (SM.isWrittenInMainFile(Loc) && Targets.count(D))
+    if (SM.isWrittenInMainFile(Loc) && CanonicalTargets.count(D))
       References.push_back({D, Loc, Roles});
     return true;
   }
 
 private:
-  llvm::SmallSet<const Decl *, 4> Targets;
+  llvm::SmallSet<const Decl *, 4> CanonicalTargets;
   std::vector<Reference> References;
   const ASTContext &AST;
 };
diff --git a/clang-tools-extra/unittests/clangd/XRefsTests.cpp b/clang-tools-extra/unittests/clangd/XRefsTests.cpp
index be4acf0977d..f6fad9112c1 100644
--- a/clang-tools-extra/unittests/clangd/XRefsTests.cpp
+++ b/clang-tools-extra/unittests/clangd/XRefsTests.cpp
@@ -1113,37 +1113,45 @@ TEST(FindReferences, WithinAST) {
   const char *Tests[] = {
       R"cpp(// Local variable
         int main() {
-          int $foo[[foo]];
-          $foo[[^foo]] = 2;
-          int test1 = $foo[[foo]];
+          int [[foo]];
+          [[^foo]] = 2;
+          int test1 = [[foo]];
         }
       )cpp",
 
       R"cpp(// Struct
         namespace ns1 {
-        struct $foo[[Foo]] {};
+        struct [[Foo]] {};
         } // namespace ns1
         int main() {
-          ns1::$foo[[Fo^o]]* Params;
+          ns1::[[Fo^o]]* Params;
+        }
+      )cpp",
+
+      R"cpp(// Forward declaration
+        class [[Foo]];
+        class [[Foo]] {}
+        int main() {
+          [[Fo^o]] foo;
         }
       )cpp",
 
       R"cpp(// Function
-        int $foo[[foo]](int) {}
+        int [[foo]](int) {}
         int main() {
-          auto *X = &$foo[[^foo]];
-          $foo[[foo]](42)
+          auto *X = &[[^foo]];
+          [[foo]](42)
         }
       )cpp",
 
       R"cpp(// Field
         struct Foo {
-          int $foo[[foo]];
-          Foo() : $foo[[foo]](0) {}
+          int [[foo]];
+          Foo() : [[foo]](0) {}
         };
         int main() {
           Foo f;
-          f.$foo[[f^oo]] = 1;
+          f.[[f^oo]] = 1;
         }
       )cpp",
 
@@ -1152,29 +1160,29 @@ TEST(FindReferences, WithinAST) {
         int Foo::[[foo]]() {}
         int main() {
           Foo f;
-          f.^foo();
+          f.[[^foo]]();
         }
       )cpp",
 
       R"cpp(// Typedef
-        typedef int $foo[[Foo]];
+        typedef int [[Foo]];
         int main() {
-          $foo[[^Foo]] bar;
+          [[^Foo]] bar;
         }
       )cpp",
 
       R"cpp(// Namespace
-        namespace $foo[[ns]] {
+        namespace [[ns]] {
         struct Foo {};
         } // namespace ns
-        int main() { $foo[[^ns]]::Foo foo; }
+        int main() { [[^ns]]::Foo foo; }
       )cpp",
   };
   for (const char *Test : Tests) {
     Annotations T(Test);
     auto AST = TestTU::withCode(T.code()).build();
     std::vector<Matcher<Location>> ExpectedLocations;
-    for (const auto &R : T.ranges("foo"))
+    for (const auto &R : T.ranges())
       ExpectedLocations.push_back(RangeIs(R));
     EXPECT_THAT(findReferences(AST, T.point()),
                 ElementsAreArray(ExpectedLocations))