[clangd] Prefer location from codegen files when merging symbols.

Summary:
For example, if an index symbol has location in a .proto file and an AST symbol
has location in a generated .proto.h file, then we prefer location in .proto
which is more meaningful to users.

Also use `mergeSymbols` to get the preferred location between AST location and index location in go-to-def.

Reviewers: sammccall

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

Tags: #clang

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

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@353708 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 103 insertions, 8 deletions

diff --git a/clangd/XRefs.cpp b/clangd/XRefs.cpp
index 3f89165b..5c97c6a6 100644
--- a/clangd/XRefs.cpp
+++ b/clangd/XRefs.cpp
@@ -10,6 +10,8 @@
 #include "Logger.h"
 #include "SourceCode.h"
 #include "URI.h"
+#include "index/Index.h"
+#include "index/Merge.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Index/IndexDataConsumer.h"
@@ -77,6 +79,32 @@ llvm::Optional<Location> toLSPLocation(const SymbolLocation &Loc,
   return LSPLoc;
 }
 
+SymbolLocation toIndexLocation(const Location &Loc, std::string &URIStorage) {
+  SymbolLocation SymLoc;
+  URIStorage = Loc.uri.uri();
+  SymLoc.FileURI = URIStorage.c_str();
+  SymLoc.Start.setLine(Loc.range.start.line);
+  SymLoc.Start.setColumn(Loc.range.start.character);
+  SymLoc.End.setLine(Loc.range.end.line);
+  SymLoc.End.setColumn(Loc.range.end.character);
+  return SymLoc;
+}
+
+// Returns the preferred location between an AST location and an index location.
+SymbolLocation getPreferredLocation(const Location &ASTLoc,
+                                    const SymbolLocation &IdxLoc) {
+  // Also use a dummy symbol for the index location so that other fields (e.g.
+  // definition) are not factored into the preferrence.
+  Symbol ASTSym, IdxSym;
+  ASTSym.ID = IdxSym.ID = SymbolID("dummy_id");
+  std::string URIStore;
+  ASTSym.CanonicalDeclaration = toIndexLocation(ASTLoc, URIStore);
+  IdxSym.CanonicalDeclaration = IdxLoc;
+  auto Merged = mergeSymbol(ASTSym, IdxSym);
+  return Merged.CanonicalDeclaration;
+}
+
+
 struct MacroDecl {
   llvm::StringRef Name;
   const MacroInfo *Info;
@@ -329,12 +357,23 @@ std::vector<LocatedSymbol> locateSymbolAt(ParsedAST &AST, Position Pos,
 
       // 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
+      if (R.Definition) { // from AST
         if (auto Loc = toLSPLocation(Sym.CanonicalDeclaration, *MainFilePath))
           R.PreferredDeclaration = *Loc;
-
-      if (!R.Definition)
+      } 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),
+                                *MainFilePath))
+            R.PreferredDeclaration = *Loc;
+        }
+      }
     });
   }
 
diff --git a/clangd/index/Merge.cpp b/clangd/index/Merge.cpp
index d967db8a..65e9b86d 100644
--- a/clangd/index/Merge.cpp
+++ b/clangd/index/Merge.cpp
@@ -9,9 +9,13 @@
 #include "Merge.h"
 #include "Logger.h"
 #include "Trace.h"
+#include "index/Index.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/raw_ostream.h"
+#include <algorithm>
+#include <iterator>
 
 namespace clang {
 namespace clangd {
@@ -114,6 +118,23 @@ void MergedIndex::refs(const RefsRequest &Req,
   });
 }
 
+// Returns true if \p L is (strictly) preferred to \p R (e.g. by file paths). If
+// neither is preferred, this returns false.
+bool prefer(const SymbolLocation &L, const SymbolLocation &R) {
+  if (!L)
+    return false;
+  if (!R)
+    return true;
+  auto HasCodeGenSuffix = [](const SymbolLocation &Loc) {
+    constexpr static const char *CodegenSuffixes[] = {".proto"};
+    return std::any_of(std::begin(CodegenSuffixes), std::end(CodegenSuffixes),
+                       [&](llvm::StringRef Suffix) {
+                         return llvm::StringRef(Loc.FileURI).endswith(Suffix);
+                       });
+  };
+  return HasCodeGenSuffix(L) && !HasCodeGenSuffix(R);
+}
+
 Symbol mergeSymbol(const Symbol &L, const Symbol &R) {
   assert(L.ID == R.ID);
   // We prefer information from TUs that saw the definition.
@@ -128,12 +149,11 @@ Symbol mergeSymbol(const Symbol &L, const Symbol &R) {
   Symbol S = PreferR ? R : L;        // The target symbol we're merging into.
   const Symbol &O = PreferR ? L : R; // The "other" less-preferred symbol.
 
-  // For each optional field, fill it from O if missing in S.
-  // (It might be missing in O too, but that's a no-op).
-  if (!S.Definition)
-    S.Definition = O.Definition;
-  if (!S.CanonicalDeclaration)
+  // Only use locations in \p O if it's (strictly) preferred.
+  if (prefer(O.CanonicalDeclaration, S.CanonicalDeclaration))
     S.CanonicalDeclaration = O.CanonicalDeclaration;
+  if (prefer(O.Definition, S.Definition))
+    S.Definition = O.Definition;
   S.References += O.References;
   if (S.Signature == "")
     S.Signature = O.Signature;
diff --git a/unittests/clangd/IndexTests.cpp b/unittests/clangd/IndexTests.cpp
index 4ca88cae..7d60ede1 100644
--- a/unittests/clangd/IndexTests.cpp
+++ b/unittests/clangd/IndexTests.cpp
@@ -253,6 +253,22 @@ TEST(MergeTest, PreferSymbolWithDefn) {
   EXPECT_EQ(M.Name, "right");
 }
 
+TEST(MergeTest, PreferSymbolLocationInCodegenFile) {
+  Symbol L, R;
+
+  L.ID = R.ID = SymbolID("hello");
+  L.CanonicalDeclaration.FileURI = "file:/x.proto.h";
+  R.CanonicalDeclaration.FileURI = "file:/x.proto";
+
+  Symbol M = mergeSymbol(L, R);
+  EXPECT_EQ(StringRef(M.CanonicalDeclaration.FileURI), "file:/x.proto");
+
+  // Prefer L if both have codegen suffix.
+  L.CanonicalDeclaration.FileURI = "file:/y.proto";
+  M = mergeSymbol(L, R);
+  EXPECT_EQ(StringRef(M.CanonicalDeclaration.FileURI), "file:/y.proto");
+}
+
 TEST(MergeIndexTest, Refs) {
   FileIndex Dyn;
   FileIndex StaticIndex;
diff --git a/unittests/clangd/XRefsTests.cpp b/unittests/clangd/XRefsTests.cpp
index acbeaf48..fea90137 100644
--- a/unittests/clangd/XRefsTests.cpp
+++ b/unittests/clangd/XRefsTests.cpp
@@ -184,6 +184,26 @@ TEST(LocateSymbol, WithIndex) {
       ElementsAre(Sym("Forward", SymbolHeader.range("forward"), Test.range())));
 }
 
+TEST(LocateSymbol, WithIndexPreferredLocation) {
+  Annotations SymbolHeader(R"cpp(
+        class $[[Proto]] {};
+      )cpp");
+  TestTU TU;
+  TU.HeaderCode = SymbolHeader.code();
+  TU.HeaderFilename = "x.proto"; // Prefer locations in codegen files.
+  auto Index = TU.index();
+
+  Annotations Test(R"cpp(// only declaration in AST.
+        // Shift to make range different.
+        class [[Proto]];
+        P^roto* create();
+      )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())));
+}
+
 TEST(LocateSymbol, All) {
   // Ranges in tests:
   //   $decl is the declaration location (if absent, no symbol is located)