[clangd] Add xref for macros to FileIndex.

Summary:
Adds macro references to the dynamic index.
Tests added.
Also exposed a new API to convert path to URI in URI.h

Reviewers: hokein

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

Tags: #clang

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

5 files changed, 70 insertions, 3 deletions

diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index fb5487689ff..179e143dc00 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "FileIndex.h"
+#include "CollectMacros.h"
 #include "Logger.h"
 #include "ParsedAST.h"
 #include "SymbolCollector.h"
@@ -32,6 +33,7 @@ namespace clangd {
 
 static SlabTuple indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
                               llvm::ArrayRef<Decl *> DeclsToIndex,
+                              const MainFileMacros *MacroRefsToIndex,
                               const CanonicalIncludes &Includes,
                               bool IsIndexMainAST) {
   SymbolCollector::Options CollectorOpts;
@@ -59,6 +61,8 @@ static SlabTuple indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
 
   SymbolCollector Collector(std::move(CollectorOpts));
   Collector.setPreprocessor(PP);
+  if (MacroRefsToIndex)
+    Collector.handleMacros(*MacroRefsToIndex);
   index::indexTopLevelDecls(AST, *PP, DeclsToIndex, Collector, IndexOpts);
 
   const auto &SM = AST.getSourceManager();
@@ -68,6 +72,7 @@ static SlabTuple indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
   auto Syms = Collector.takeSymbols();
   auto Refs = Collector.takeRefs();
   auto Relations = Collector.takeRelations();
+
   vlog("index AST for {0} (main={1}): \n"
        "  symbol slab: {2} symbols, {3} bytes\n"
        "  ref slab: {4} symbols, {5} refs, {6} bytes\n"
@@ -80,7 +85,8 @@ static SlabTuple indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
 
 SlabTuple indexMainDecls(ParsedAST &AST) {
   return indexSymbols(AST.getASTContext(), AST.getPreprocessorPtr(),
-                      AST.getLocalTopLevelDecls(), AST.getCanonicalIncludes(),
+                      AST.getLocalTopLevelDecls(), &AST.getMacros(),
+                      AST.getCanonicalIncludes(),
                       /*IsIndexMainAST=*/true);
 }
 
@@ -89,7 +95,8 @@ SlabTuple indexHeaderSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
   std::vector<Decl *> DeclsToIndex(
       AST.getTranslationUnitDecl()->decls().begin(),
       AST.getTranslationUnitDecl()->decls().end());
-  return indexSymbols(AST, std::move(PP), DeclsToIndex, Includes,
+  return indexSymbols(AST, std::move(PP), DeclsToIndex,
+                      /*MainFileMacros=*/nullptr, Includes,
                       /*IsIndexMainAST=*/false);
 }
 
diff --git a/clang-tools-extra/clangd/index/FileIndex.h b/clang-tools-extra/clangd/index/FileIndex.h
index 411130c5954..50f297cc198 100644
--- a/clang-tools-extra/clangd/index/FileIndex.h
+++ b/clang-tools-extra/clangd/index/FileIndex.h
@@ -140,7 +140,7 @@ using SlabTuple = std::tuple<SymbolSlab, RefSlab, RelationSlab>;
 /// Exposed to assist in unit tests.
 SlabTuple indexMainDecls(ParsedAST &AST);
 
-/// Idex declarations from \p AST and macros from \p PP that are declared in
+/// Index declarations from \p AST and macros from \p PP that are declared in
 /// included headers.
 SlabTuple indexHeaderSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
                              const CanonicalIncludes &Includes);
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.cpp b/clang-tools-extra/clangd/index/SymbolCollector.cpp
index c0b2b52a43c..48c26bae5c2 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ b/clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -346,6 +346,29 @@ bool SymbolCollector::handleDeclOccurrence(
   return true;
 }
 
+void SymbolCollector::handleMacros(const MainFileMacros &MacroRefsToIndex) {
+  assert(PP.get());
+  const auto &SM = PP->getSourceManager();
+  const auto *MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
+  assert(MainFileEntry);
+
+  const auto MainFileURI = toURI(SM, MainFileEntry->getName(), Opts);
+  // Add macro references.
+  for (const auto &IDToRefs : MacroRefsToIndex.MacroRefs) {
+    for (const auto &Range : IDToRefs.second) {
+      Ref R;
+      R.Location.Start.setLine(Range.start.line);
+      R.Location.Start.setColumn(Range.start.character);
+      R.Location.End.setLine(Range.end.line);
+      R.Location.End.setColumn(Range.end.character);
+      R.Location.FileURI = MainFileURI.c_str();
+      // FIXME: Add correct RefKind information to MainFileMacros.
+      R.Kind = RefKind::Reference;
+      Refs.insert(IDToRefs.first, R);
+    }
+  }
+}
+
 bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
                                             const MacroInfo *MI,
                                             index::SymbolRoleSet Roles,
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.h b/clang-tools-extra/clangd/index/SymbolCollector.h
index 7643b1f8286..f66a71c2d59 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.h
+++ b/clang-tools-extra/clangd/index/SymbolCollector.h
@@ -9,6 +9,7 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOL_COLLECTOR_H
 
 #include "CanonicalIncludes.h"
+#include "CollectMacros.h"
 #include "Index.h"
 #include "SymbolOrigin.h"
 #include "clang/AST/ASTContext.h"
@@ -108,6 +109,8 @@ public:
                              index::SymbolRoleSet Roles,
                              SourceLocation Loc) override;
 
+  void handleMacros(const MainFileMacros &MacroRefsToIndex);
+
   SymbolSlab takeSymbols() { return std::move(Symbols).build(); }
   RefSlab takeRefs() { return std::move(Refs).build(); }
   RelationSlab takeRelations() { return std::move(Relations).build(); }
diff --git a/clang-tools-extra/clangd/unittests/FileIndexTests.cpp b/clang-tools-extra/clangd/unittests/FileIndexTests.cpp
index bc07217835b..9418b6897cb 100644
--- a/clang-tools-extra/clangd/unittests/FileIndexTests.cpp
+++ b/clang-tools-extra/clangd/unittests/FileIndexTests.cpp
@@ -345,6 +345,40 @@ TEST(FileIndexTest, Refs) {
                              FileURI("unittest:///test2.cc"))}));
 }
 
+TEST(FileIndexTest, MacroRefs) {
+  Annotations HeaderCode(R"cpp(
+    #define $def1[[HEADER_MACRO]](X) (X+1)
+  )cpp");
+  Annotations MainCode(R"cpp(
+  #define $def2[[MAINFILE_MACRO]](X) (X+1)
+  void f() {
+    int a = $ref1[[HEADER_MACRO]](2);
+    int b = $ref2[[MAINFILE_MACRO]](1);
+  }
+  )cpp");
+
+  FileIndex Index;
+  // Add test.cc
+  TestTU Test;
+  Test.HeaderCode = HeaderCode.code();
+  Test.Code = MainCode.code();
+  Test.Filename = "test.cc";
+  auto AST = Test.build();
+  Index.updateMain(Test.Filename, AST);
+
+  auto HeaderMacro = findSymbol(Test.headerSymbols(), "HEADER_MACRO");
+  EXPECT_THAT(getRefs(Index, HeaderMacro.ID),
+              RefsAre({AllOf(RefRange(MainCode.range("ref1")),
+                             FileURI("unittest:///test.cc"))}));
+
+  auto MainFileMacro = findSymbol(Test.headerSymbols(), "MAINFILE_MACRO");
+  EXPECT_THAT(getRefs(Index, MainFileMacro.ID),
+              RefsAre({AllOf(RefRange(MainCode.range("def2")),
+                             FileURI("unittest:///test.cc")),
+                       AllOf(RefRange(MainCode.range("ref2")),
+                             FileURI("unittest:///test.cc"))}));
+}
+
 TEST(FileIndexTest, CollectMacros) {
   FileIndex M;
   update(M, "f", "#define CLANGD 1");