[clangd] Index refs to main-file symbols as well

Summary: This will be needed to support call hierarchy

Reviewers: kadircet

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

Tags: #clang

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

Author: HighCommander4

clang-tools-extra/clangd/ClangdServer.cpp

@@ -173,7 +173,8 @@ ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            Callbacks *Callbacks)
     : ConfigProvider(Opts.ConfigProvider), TFS(TFS),
       DynamicIdx(Opts.BuildDynamicSymbolIndex
-                     ? new FileIndex(Opts.HeavyweightDynamicSymbolIndex)
+                     ? new FileIndex(Opts.HeavyweightDynamicSymbolIndex,
+                                     Opts.CollectMainFileRefs)
                      : nullptr),
       GetClangTidyOptions(Opts.GetClangTidyOptions),
       SuggestMissingIncludes(Opts.SuggestMissingIncludes),

clang-tools-extra/clangd/ClangdServer.h

@@ -111,6 +111,9 @@ class ClangdServer {
     /// on background threads. The index is stored in the project root.
     bool BackgroundIndex = false;
 
+    /// Store refs to main-file symbols in the index.
+    bool CollectMainFileRefs = false;
+
     /// If set, use this index to augment code completion results.
     SymbolIndex *StaticIndex = nullptr;
 

clang-tools-extra/clangd/index/Background.cpp

@@ -95,6 +95,7 @@ BackgroundIndex::BackgroundIndex(
     BackgroundIndexStorage::Factory IndexStorageFactory, Options Opts)
     : SwapIndex(std::make_unique<MemIndex>()), TFS(TFS), CDB(CDB),
       ContextProvider(std::move(Opts.ContextProvider)),
+      CollectMainFileRefs(Opts.CollectMainFileRefs),
       Rebuilder(this, &IndexedSymbols, Opts.ThreadPoolSize),
       IndexStorageFactory(std::move(IndexStorageFactory)),
       Queue(std::move(Opts.OnProgress)),
@@ -301,6 +302,7 @@ llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
       return false; // Skip files that haven't changed, without errors.
     return true;
   };
+  IndexOpts.CollectMainFileRefs = CollectMainFileRefs;
 
   IndexFileIn Index;
   auto Action = createStaticIndexingAction(

clang-tools-extra/clangd/index/Background.h

@@ -137,6 +137,8 @@ class BackgroundIndex : public SwapIndex {
     // file. Called with the empty string for other tasks.
     // (When called, the context from BackgroundIndex construction is active).
     std::function<Context(PathRef)> ContextProvider = nullptr;
+    // Whether to collect references to main-file-only symbols.
+    bool CollectMainFileRefs = false;
   };
 
   /// Creates a new background index and starts its threads.
@@ -188,6 +190,7 @@ class BackgroundIndex : public SwapIndex {
   const ThreadsafeFS &TFS;
   const GlobalCompilationDatabase &CDB;
   std::function<Context(PathRef)> ContextProvider;
+  bool CollectMainFileRefs;
 
   llvm::Error index(tooling::CompileCommand);
 

clang-tools-extra/clangd/index/FileIndex.cpp

@@ -47,12 +47,13 @@ SlabTuple indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
                        llvm::ArrayRef<Decl *> DeclsToIndex,
                        const MainFileMacros *MacroRefsToIndex,
                        const CanonicalIncludes &Includes, bool IsIndexMainAST,
-                       llvm::StringRef Version) {
+                       llvm::StringRef Version, bool CollectMainFileRefs) {
   SymbolCollector::Options CollectorOpts;
   CollectorOpts.CollectIncludePath = true;
   CollectorOpts.Includes = &Includes;
   CollectorOpts.CountReferences = false;
   CollectorOpts.Origin = SymbolOrigin::Dynamic;
+  CollectorOpts.CollectMainFileRefs = CollectMainFileRefs;
 
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
@@ -205,11 +206,11 @@ FileShardedIndex::getShard(llvm::StringRef Uri) const {
   return std::move(IF);
 }
 
-SlabTuple indexMainDecls(ParsedAST &AST) {
-  return indexSymbols(AST.getASTContext(), AST.getPreprocessorPtr(),
-                      AST.getLocalTopLevelDecls(), &AST.getMacros(),
-                      AST.getCanonicalIncludes(),
-                      /*IsIndexMainAST=*/true, AST.version());
+SlabTuple indexMainDecls(ParsedAST &AST, bool CollectMainFileRefs) {
+  return indexSymbols(
+      AST.getASTContext(), AST.getPreprocessorPtr(),
+      AST.getLocalTopLevelDecls(), &AST.getMacros(), AST.getCanonicalIncludes(),
+      /*IsIndexMainAST=*/true, AST.version(), CollectMainFileRefs);
 }
 
 SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
@@ -220,7 +221,8 @@ SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
       AST.getTranslationUnitDecl()->decls().end());
   return indexSymbols(AST, std::move(PP), DeclsToIndex,
                       /*MainFileMacros=*/nullptr, Includes,
-                      /*IsIndexMainAST=*/false, Version);
+                      /*IsIndexMainAST=*/false, Version,
+                      /*CollectMainFileRefs=*/false);
 }
 
 void FileSymbols::update(llvm::StringRef Key,
@@ -371,8 +373,9 @@ FileSymbols::buildIndex(IndexType Type, DuplicateHandling DuplicateHandle,
   llvm_unreachable("Unknown clangd::IndexType");
 }
 
-FileIndex::FileIndex(bool UseDex)
+FileIndex::FileIndex(bool UseDex, bool CollectMainFileRefs)
     : MergedIndex(&MainFileIndex, &PreambleIndex), UseDex(UseDex),
+      CollectMainFileRefs(CollectMainFileRefs),
       PreambleIndex(std::make_unique<MemIndex>()),
       MainFileIndex(std::make_unique<MemIndex>()) {}
 
@@ -415,7 +418,7 @@ void FileIndex::updatePreamble(PathRef Path, llvm::StringRef Version,
 }
 
 void FileIndex::updateMain(PathRef Path, ParsedAST &AST) {
-  auto Contents = indexMainDecls(AST);
+  auto Contents = indexMainDecls(AST, CollectMainFileRefs);
   MainFileSymbols.update(
       Path, std::make_unique<SymbolSlab>(std::move(std::get<0>(Contents))),
       std::make_unique<RefSlab>(std::move(std::get<1>(Contents))),

clang-tools-extra/clangd/index/FileIndex.h

@@ -104,7 +104,7 @@ class FileSymbols {
 /// FIXME: Expose an interface to remove files that are closed.
 class FileIndex : public MergedIndex {
 public:
-  FileIndex(bool UseDex = true);
+  FileIndex(bool UseDex = true, bool CollectMainFileRefs = false);
 
   /// Update preamble symbols of file \p Path with all declarations in \p AST
   /// and macros in \p PP.
@@ -118,6 +118,7 @@ class FileIndex : public MergedIndex {
 
 private:
   bool UseDex; // FIXME: this should be always on.
+  bool CollectMainFileRefs;
 
   // Contains information from each file's preamble only. Symbols and relations
   // are sharded per declaration file to deduplicate multiple symbols and reduce
@@ -152,7 +153,7 @@ using SlabTuple = std::tuple<SymbolSlab, RefSlab, RelationSlab>;
 /// Retrieves symbols and refs of local top level decls in \p AST (i.e.
 /// `AST.getLocalTopLevelDecls()`).
 /// Exposed to assist in unit tests.
-SlabTuple indexMainDecls(ParsedAST &AST);
+SlabTuple indexMainDecls(ParsedAST &AST, bool CollectMainFileRefs = false);
 
 /// Index declarations from \p AST and macros from \p PP that are declared in
 /// included headers.

clang-tools-extra/clangd/index/SymbolCollector.cpp

@@ -334,12 +334,13 @@ bool SymbolCollector::handleDeclOccurrence(
   if (IsOnlyRef && !CollectRef)
     return true;
 
-  // Do not store references to main-file symbols.
   // Unlike other fields, e.g. Symbols (which use spelling locations), we use
   // file locations for references (as it aligns the behavior of clangd's
   // AST-based xref).
   // FIXME: we should try to use the file locations for other fields.
-  if (CollectRef && (!IsMainFileOnly || ND->isExternallyVisible()) &&
+  if (CollectRef &&
+      (!IsMainFileOnly || Opts.CollectMainFileRefs ||
+       ND->isExternallyVisible()) &&
       !isa<NamespaceDecl>(ND) &&
       (Opts.RefsInHeaders ||
        SM.getFileID(SM.getFileLoc(Loc)) == SM.getMainFileID()))

clang-tools-extra/clangd/index/SymbolCollector.h

@@ -78,6 +78,8 @@ class SymbolCollector : public index::IndexDataConsumer {
     /// Collect symbols local to main-files, such as static functions
     /// and symbols inside an anonymous namespace.
     bool CollectMainFileSymbols = true;
+    /// Collect references to main-file symbols.
+    bool CollectMainFileRefs = false;
     /// If set to true, SymbolCollector will collect doc for all symbols.
     /// Note that documents of symbols being indexed for completion will always
     /// be collected regardless of this option.

clang-tools-extra/clangd/tool/ClangdMain.cpp

@@ -450,6 +450,13 @@ opt<bool> EnableConfig{
     init(true),
 };
 
+opt<bool> CollectMainFileRefs{
+    "collect-main-file-refs",
+    cat(Misc),
+    desc("Store references to main-file-only symbols in the index"),
+    init(false),
+};
+
 #if CLANGD_ENABLE_REMOTE
 opt<std::string> RemoteIndexAddress{
     "remote-index-address",
@@ -682,6 +689,7 @@ clangd accepts flags on the commandline, and in the CLANGD_FLAGS environment var
   if (!ResourceDir.empty())
     Opts.ResourceDir = ResourceDir;
   Opts.BuildDynamicSymbolIndex = EnableIndex;
+  Opts.CollectMainFileRefs = CollectMainFileRefs;
   std::unique_ptr<SymbolIndex> StaticIdx;
   std::future<void> AsyncIndexLoad; // Block exit while loading the index.
   if (EnableIndex && !IndexFile.empty()) {

clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp

@@ -229,6 +229,61 @@ TEST_F(BackgroundIndexTest, IndexTwoFiles) {
                        FileURI("unittest:///root/B.cc")}));
 }
 
+TEST_F(BackgroundIndexTest, MainFileRefs) {
+  MockFS FS;
+  FS.Files[testPath("root/A.h")] = R"cpp(
+      void header_sym();
+      )cpp";
+  FS.Files[testPath("root/A.cc")] =
+      "#include \"A.h\"\nstatic void main_sym() { (void)header_sym; }";
+
+  // Check the behaviour with CollectMainFileRefs = false (the default).
+  {
+    llvm::StringMap<std::string> Storage;
+    size_t CacheHits = 0;
+    MemoryShardStorage MSS(Storage, CacheHits);
+    OverlayCDB CDB(/*Base=*/nullptr);
+    BackgroundIndex Idx(FS, CDB, [&](llvm::StringRef) { return &MSS; },
+                        /*Opts=*/{});
+
+    tooling::CompileCommand Cmd;
+    Cmd.Filename = testPath("root/A.cc");
+    Cmd.Directory = testPath("root");
+    Cmd.CommandLine = {"clang++", testPath("root/A.cc")};
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
+
+    ASSERT_TRUE(Idx.blockUntilIdleForTest());
+    EXPECT_THAT(
+        runFuzzyFind(Idx, ""),
+        UnorderedElementsAre(AllOf(Named("header_sym"), NumReferences(1U)),
+                             AllOf(Named("main_sym"), NumReferences(0U))));
+  }
+
+  // Check the behaviour with CollectMainFileRefs = true.
+  {
+    llvm::StringMap<std::string> Storage;
+    size_t CacheHits = 0;
+    MemoryShardStorage MSS(Storage, CacheHits);
+    OverlayCDB CDB(/*Base=*/nullptr);
+    BackgroundIndex::Options Opts;
+    Opts.CollectMainFileRefs = true;
+    BackgroundIndex Idx(
+        FS, CDB, [&](llvm::StringRef) { return &MSS; }, Opts);
+
+    tooling::CompileCommand Cmd;
+    Cmd.Filename = testPath("root/A.cc");
+    Cmd.Directory = testPath("root");
+    Cmd.CommandLine = {"clang++", testPath("root/A.cc")};
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
+
+    ASSERT_TRUE(Idx.blockUntilIdleForTest());
+    EXPECT_THAT(
+        runFuzzyFind(Idx, ""),
+        UnorderedElementsAre(AllOf(Named("header_sym"), NumReferences(1U)),
+                             AllOf(Named("main_sym"), NumReferences(1U))));
+  }
+}
+
 TEST_F(BackgroundIndexTest, ShardStorageTest) {
   MockFS FS;
   FS.Files[testPath("root/A.h")] = R"cpp(

clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp

@@ -714,7 +714,6 @@ TEST_F(SymbolCollectorTest, Refs) {
   EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(Symbols, "NS").ID, _))));
   EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "MACRO").ID,
                                   HaveRanges(Main.ranges("macro")))));
-  // Symbols *only* in the main file:
   // - (a, b) externally visible and should have refs.
   // - (c, FUNC) externally invisible and had no refs collected.
   auto MainSymbols =
@@ -723,6 +722,20 @@ TEST_F(SymbolCollectorTest, Refs) {
   EXPECT_THAT(Refs, Contains(Pair(findSymbol(MainSymbols, "b").ID, _)));
   EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(MainSymbols, "c").ID, _))));
   EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(MainSymbols, "FUNC").ID, _))));
+
+  // Run the collector again with CollectMainFileRefs = true.
+  // We need to recreate InMemoryFileSystem because runSymbolCollector()
+  // calls MemoryBuffer::getMemBuffer(), which makes the buffers unusable
+  // after runSymbolCollector() exits.
+  InMemoryFileSystem = new llvm::vfs::InMemoryFileSystem();
+  CollectorOpts.CollectMainFileRefs = true;
+  runSymbolCollector(Header.code(),
+                     (Main.code() + SymbolsOnlyInMainCode.code()).str());
+  EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "a").ID, _)));
+  EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "b").ID, _)));
+  EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "c").ID, _)));
+  // However, references to main-file macros are not collected.
+  EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(Symbols, "FUNC").ID, _))));
 }
 
 TEST_F(SymbolCollectorTest, MacroRefInHeader) {
@@ -908,8 +921,9 @@ TEST_F(SymbolCollectorTest, HeaderAsMainFile) {
     $Foo[[Foo]] fo;
   }
   )");
-  // The main file is normal .cpp file, we should collect the refs
-  // for externally visible symbols.
+  // We should collect refs to main-file symbols in all cases:
+
+  // 1. The main file is normal .cpp file.
   TestFileName = testPath("foo.cpp");
   runSymbolCollector("", Header.code());
   EXPECT_THAT(Refs,
@@ -918,7 +932,7 @@ TEST_F(SymbolCollectorTest, HeaderAsMainFile) {
                                    Pair(findSymbol(Symbols, "Func").ID,
                                         HaveRanges(Header.ranges("Func")))));
 
-  // Run the .h file as main file, we should collect the refs.
+  // 2. Run the .h file as main file.
   TestFileName = testPath("foo.h");
   runSymbolCollector("", Header.code(),
                      /*ExtraArgs=*/{"-xobjective-c++-header"});
@@ -929,8 +943,7 @@ TEST_F(SymbolCollectorTest, HeaderAsMainFile) {
                                    Pair(findSymbol(Symbols, "Func").ID,
                                         HaveRanges(Header.ranges("Func")))));
 
-  // Run the .hh file as main file (without "-x c++-header"), we should collect
-  // the refs as well.
+  // 3. Run the .hh file as main file (without "-x c++-header").
   TestFileName = testPath("foo.hh");
   runSymbolCollector("", Header.code());
   EXPECT_THAT(Symbols, UnorderedElementsAre(QName("Foo"), QName("Func")));