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

[ioeric]

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.

[sammccall]

clangd/XRefs.cpp

261

nit: move this and toIndexLocation together?

364

a little unclear what "we only do this" refers to - maybe missing a first line of the comment here.

// Use merge logic to choose AST or index declaration.

clangd/index/Merge.cpp

155

I think it might be clearer what this is doing, and more similar with surrounding code, if the function returns a bool:

bool prefer(const SymbolLocation &Candidate, const SymbolLocation &Baseline);
...
if (!S.Definition || prefer(O.Definition, S.Definition))
  S.Definition = O.Definition;

up to you

[ioeric]

• Address review comments.

clangd/XRefs.cpp

 //===--- XRefs.cpp -----------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #include "XRefs.h"
 #include "AST.h"
 #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"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Index/USRGeneration.h"
 #include "llvm/Support/Path.h"
 
 namespace clang {
llvm::Optional<Location> toLSPLocation(const SymbolLocation &Loc,
   LSPLoc.range.start.line = Loc.Start.line();
   LSPLoc.range.start.character = Loc.Start.column();
   LSPLoc.range.end.line = Loc.End.line();
   LSPLoc.range.end.character = Loc.End.column();
   logIfOverflow(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;
+}
+
 struct MacroDecl {
   llvm::StringRef Name;
   const MacroInfo *Info;
 };
 
 struct DeclInfo {
   const Decl *D;
   // Indicates the declaration is referenced by an explicit AST node.
if (!FilePath) {
     log("failed to get path!");
     return None;
   }
   Location L;
   L.uri = URIForFile::canonicalize(*FilePath, TUPath);
   L.range = getTokenRange(AST, TokLoc);
   return L;
 }
+// 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;
+}
 
 } // namespace
 
 std::vector<LocatedSymbol> locateSymbolAt(ParsedAST &AST, Position Pos,
                                           const SymbolIndex *Index) {
   const auto &SM = AST.getASTContext().getSourceManager();
   auto MainFilePath =
       getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
if (Index && !ResultIndex.empty()) {
     LookupRequest QueryRequest;
     for (auto It : ResultIndex)
       QueryRequest.IDs.insert(It.first);
     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
+      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) {
+          // 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;
+        }
+      }
     });
   }
 
   return Result;
 }
 
 namespace {
 

clangd/index/Merge.cpp

 //===--- Merge.cpp -----------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #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 {
 
 // FIXME: Deleted symbols in dirty files are still returned (from Static).
 //        To identify these eliminate these, we should:
 //          - find the generating file from each Symbol which is Static-only
 //          - ask Dynamic if it has that file (needs new SymbolIndex method)
void MergedIndex::refs(const RefsRequest &Req,
   Static->refs(Req, [&](const Ref &O) {
     if (Remaining > 0 && !DynamicIndexFileURIs.count(O.Location.FileURI)) {
       --Remaining;
       Callback(O);
     }
   });
 }
 
+// Returns the preferred location e.g. by file path. \p L is returned if
+// neither is preferred.
+static SymbolLocation getPreferredLocation(const SymbolLocation &L,
+                                           const SymbolLocation &R) {
+  if (!L)
+    return R;
+  if (!R)
+    return L;
+  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(R) && !HasCodeGenSuffix(L)) ? R : L;
+}
+
 Symbol mergeSymbol(const Symbol &L, const Symbol &R) {
   assert(L.ID == R.ID);
   // We prefer information from TUs that saw the definition.
   // Classes: this is the def itself. Functions: hopefully the header decl.
   // If both did (or both didn't), continue to prefer L over R.
   bool PreferR = R.Definition && !L.Definition;
   // Merge include headers only if both have definitions or both have no
   // definition; otherwise, only accumulate references of common includes.
   assert(L.Definition.FileURI && R.Definition.FileURI);
   bool MergeIncludes =
       bool(*L.Definition.FileURI) == bool(*R.Definition.FileURI);
   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)
-    S.CanonicalDeclaration = O.CanonicalDeclaration;
+  S.CanonicalDeclaration =
+      getPreferredLocation(S.CanonicalDeclaration, O.CanonicalDeclaration);
+  S.Definition = getPreferredLocation(S.Definition, O.Definition);
   S.References += O.References;
   if (S.Signature == "")
     S.Signature = O.Signature;
   if (S.CompletionSnippetSuffix == "")
     S.CompletionSnippetSuffix = O.CompletionSnippetSuffix;
   if (S.Documentation == "")
     S.Documentation = O.Documentation;
   if (S.ReturnType == "")

unittests/clangd/IndexTests.cpp

TEST(MergeTest, PreferSymbolWithDefn) {
 
   R.Definition.FileURI = "file:/right.cpp"; // Now right will be favored.
   M = mergeSymbol(L, R);
   EXPECT_EQ(StringRef(M.CanonicalDeclaration.FileURI), "file:/right.h");
   EXPECT_EQ(StringRef(M.Definition.FileURI), "file:/right.cpp");
   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;
   MergedIndex Merge(&Dyn, &StaticIndex);
 
   const char *HeaderCode = "class Foo;";
   auto HeaderSymbols = TestTU::withHeaderCode("class Foo;").headerSymbols();
   auto Foo = findSymbol(HeaderSymbols, "Foo");

unittests/clangd/XRefsTests.cpp

Test = Annotations(R"cpp(// defintion in AST.
         class [[Forward]] {};
         F^orward create();
       )cpp");
   EXPECT_THAT(
       LocateWithIndex(Test),
       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)
   //   $def is the definition location (if absent, symbol has no definition)
   //   unnamed range becomes both $decl and $def.
   const char *Tests[] = {
       R"cpp(// Local variable
         int main() {