[clangd] Add SymbolID to LocatedSymbol.

This is useful for running in batch mode.
Getting the SymbolID from via getSymbolInfo may give SymbolID
of a symbol different from that located by LocateSymbolAt (they
have different semantics of choosing the symbol.)

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

3 files changed, 31 insertions, 2 deletions

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index e4c8db24eb20..bb51b4a48d09 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -257,6 +257,7 @@ locateMacroReferent(const syntax::Token &TouchedIdentifier, ParsedAST &AST,
       Macro.Name = std::string(M->Name);
       Macro.PreferredDeclaration = *Loc;
       Macro.Definition = Loc;
+      Macro.ID = getSymbolID(M->Name, M->Info, AST.getSourceManager());
       return Macro;
     }
   }
@@ -361,6 +362,7 @@ locateASTReferent(SourceLocation CurLoc, const syntax::Token *TouchedIdentifier,
     Result.emplace_back();
     Result.back().Name = printName(AST.getASTContext(), *D);
     Result.back().PreferredDeclaration = *Loc;
+    Result.back().ID = getSymbolID(D);
     if (const NamedDecl *Def = getDefinition(D))
       Result.back().Definition = makeLocation(
           AST.getASTContext(), nameLocation(*Def, SM), MainFilePath);
@@ -516,6 +518,7 @@ std::vector<LocatedSymbol> locateSymbolForType(const ParsedAST &AST,
     Results.emplace_back();
     Results.back().Name = printName(ASTContext, *D);
     Results.back().PreferredDeclaration = *Loc;
+    Results.back().ID = getSymbolID(D);
     if (const NamedDecl *Def = getDefinition(D))
       Results.back().Definition =
           makeLocation(ASTContext, nameLocation(*Def, SM), *MainFilePath);
@@ -605,6 +608,7 @@ locateSymbolTextually(const SpelledWord &Word, ParsedAST &AST,
     LocatedSymbol Located;
     Located.PreferredDeclaration = *MaybeDeclLoc;
     Located.Name = (Sym.Name + Sym.TemplateSpecializationArgs).str();
+    Located.ID = Sym.ID;
     if (Sym.Definition) {
       auto MaybeDefLoc = indexToLSPLocation(Sym.Definition, MainFilePath);
       if (!MaybeDefLoc) {
diff --git a/clang-tools-extra/clangd/XRefs.h b/clang-tools-extra/clangd/XRefs.h
index 3f69106611dc..aa0eaa7454c0 100644
--- a/clang-tools-extra/clangd/XRefs.h
+++ b/clang-tools-extra/clangd/XRefs.h
@@ -16,6 +16,7 @@
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "index/Index.h"
+#include "index/SymbolID.h"
 #include "index/SymbolLocation.h"
 #include "support/Path.h"
 #include "clang/AST/ASTTypeTraits.h"
@@ -47,6 +48,8 @@ struct LocatedSymbol {
   Location PreferredDeclaration;
   // Where the symbol is defined, if known. May equal PreferredDeclaration.
   llvm::Optional<Location> Definition;
+  // SymbolID of the located symbol if available.
+  SymbolID ID;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const LocatedSymbol &);
 /// Get definition of symbol at a specified \p Pos.
diff --git a/clang-tools-extra/clangd/unittests/XRefsTests.cpp b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
index 2fbf1f98db8a..8c37532507d4 100644
--- a/clang-tools-extra/clangd/unittests/XRefsTests.cpp
+++ b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
@@ -306,6 +306,7 @@ MATCHER_P(Sym, Name, "") { return arg.Name == Name; }
 
 MATCHER_P(RangeIs, R, "") { return arg.Loc.range == R; }
 MATCHER_P(AttrsAre, A, "") { return arg.Attributes == A; }
+MATCHER_P(HasID, ID, "") { return arg.ID == ID; }
 
 TEST(LocateSymbol, WithIndex) {
   Annotations SymbolHeader(R"cpp(
@@ -919,6 +920,7 @@ TEST(LocateSymbol, All) {
     } else {
       ASSERT_THAT(Results, ::testing::SizeIs(1)) << Test;
       EXPECT_EQ(Results[0].PreferredDeclaration.range, *WantDecl) << Test;
+      EXPECT_TRUE(Results[0].ID) << Test;
       llvm::Optional<Range> GotDef;
       if (Results[0].Definition)
         GotDef = Results[0].Definition->range;
@@ -926,6 +928,24 @@ TEST(LocateSymbol, All) {
     }
   }
 }
+TEST(LocateSymbol, ValidSymbolID) {
+  auto T = Annotations(R"cpp(
+    #define MACRO(x, y) ((x) + (y))
+    int add(int x, int y) { return $MACRO^MACRO(x, y); }
+    int sum = $add^add(1, 2);
+  )cpp");
+
+  TestTU TU = TestTU::withCode(T.code());
+  auto AST = TU.build();
+  auto Index = TU.index();
+  EXPECT_THAT(locateSymbolAt(AST, T.point("add"), Index.get()),
+              ElementsAre(AllOf(Sym("add"),
+                                HasID(getSymbolID(&findDecl(AST, "add"))))));
+  EXPECT_THAT(
+      locateSymbolAt(AST, T.point("MACRO"), Index.get()),
+      ElementsAre(AllOf(Sym("MACRO"),
+                        HasID(findSymbol(TU.headerSymbols(), "MACRO").ID))));
+}
 
 TEST(LocateSymbol, AllMulti) {
   // Ranges in tests:
@@ -1072,8 +1092,10 @@ TEST(LocateSymbol, TextualSmoke) {
   auto TU = TestTU::withCode(T.code());
   auto AST = TU.build();
   auto Index = TU.index();
-  EXPECT_THAT(locateSymbolAt(AST, T.point(), Index.get()),
-              ElementsAre(Sym("MyClass", T.range(), T.range())));
+  EXPECT_THAT(
+      locateSymbolAt(AST, T.point(), Index.get()),
+      ElementsAre(AllOf(Sym("MyClass", T.range(), T.range()),
+                        HasID(getSymbolID(&findDecl(AST, "MyClass"))))));
 }
 
 TEST(LocateSymbol, Textual) {