[clangd] Also perform merging for symbol definitions

Summary:
clangd currently prefers declarations from codegen files. This patch
implements that behavior for definition locations. If we have definiton
locations both coming from AST and index, clangd will perform a merging to show
the codegen file if that's the case.

Reviewers: sammccall

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

Tags: #clang

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

llvm-svn: 359874

2 files changed, 33 insertions, 17 deletions

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index 6f47618fb08..c51631ad193 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -346,24 +346,27 @@ std::vector<LocatedSymbol> locateSymbolAt(ParsedAST &AST, Position Pos,
     Index->lookup(QueryRequest, [&](const Symbol &Sym) {
       auto &R = Result[ResultIndex.lookup(Sym.ID)];
 
-      // Special case: if the AST yielded a definition, then it may not be
-      // the right *declaration*. Prefer the one from the index.
       if (R.Definition) { // from AST
+        // Special case: if the AST yielded a definition, then it may not be
+        // the right *declaration*. Prefer the one from the index.
         if (auto Loc = toLSPLocation(Sym.CanonicalDeclaration, *MainFilePath))
           R.PreferredDeclaration = *Loc;
+
+        // We might still prefer the definition from the index, e.g. for
+        // generated symbols.
+        if (auto Loc = toLSPLocation(
+                getPreferredLocation(*R.Definition, Sym.Definition, Scratch),
+                *MainFilePath))
+          R.Definition = *Loc;
       } else {
         R.Definition = toLSPLocation(Sym.Definition, *MainFilePath);
 
-        if (Sym.CanonicalDeclaration) {
-          // Use merge logic to choose AST or index declaration.
-          // We only do this for declarations as definitions from AST
-          // is generally preferred (e.g. definitions in main file).
-          if (auto Loc = toLSPLocation(
-                  getPreferredLocation(R.PreferredDeclaration,
-                                       Sym.CanonicalDeclaration, Scratch),
-                  *MainFilePath))
-            R.PreferredDeclaration = *Loc;
-        }
+        // Use merge logic to choose AST or index declaration.
+        if (auto Loc = toLSPLocation(
+                getPreferredLocation(R.PreferredDeclaration,
+                                     Sym.CanonicalDeclaration, Scratch),
+                *MainFilePath))
+          R.PreferredDeclaration = *Loc;
       }
     });
   }
diff --git a/clang-tools-extra/clangd/unittests/XRefsTests.cpp b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
index a9fb5898435..016753ca55f 100644
--- a/clang-tools-extra/clangd/unittests/XRefsTests.cpp
+++ b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
@@ -186,7 +186,8 @@ TEST(LocateSymbol, WithIndex) {
 
 TEST(LocateSymbol, WithIndexPreferredLocation) {
   Annotations SymbolHeader(R"cpp(
-        class $[[Proto]] {};
+        class $p[[Proto]] {};
+        void $f[[func]]() {};
       )cpp");
   TestTU TU;
   TU.HeaderCode = SymbolHeader.code();
@@ -195,13 +196,25 @@ TEST(LocateSymbol, WithIndexPreferredLocation) {
 
   Annotations Test(R"cpp(// only declaration in AST.
         // Shift to make range different.
-        class [[Proto]];
-        P^roto* create();
+        class Proto;
+        void func() {}
+        P$p^roto* create() {
+          fu$f^nc();
+          return nullptr;
+        }
       )cpp");
 
   auto AST = TestTU::withCode(Test.code()).build();
-  auto Locs = clangd::locateSymbolAt(AST, Test.point(), Index.get());
-  EXPECT_THAT(Locs, ElementsAre(Sym("Proto", SymbolHeader.range())));
+  {
+    auto Locs = clangd::locateSymbolAt(AST, Test.point("p"), Index.get());
+    auto CodeGenLoc = SymbolHeader.range("p");
+    EXPECT_THAT(Locs, ElementsAre(Sym("Proto", CodeGenLoc, CodeGenLoc)));
+  }
+  {
+    auto Locs = clangd::locateSymbolAt(AST, Test.point("f"), Index.get());
+    auto CodeGenLoc = SymbolHeader.range("f");
+    EXPECT_THAT(Locs, ElementsAre(Sym("func", CodeGenLoc, CodeGenLoc)));
+  }
 }
 
 TEST(LocateSymbol, All) {