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

[nridge]

Summary

This will be needed to support call hierarchy

[nridge]

(As per discussion in https://github.com/clangd/clangd/issues/162#issuecomment-648923948 onwards.)

[kadircet]

Thanks a lot for doing this Nathan and sorry for not replying with a concrete decision on the issue tracker, we were still discussing it internally, as 7% increase for allowing multi-level call hierarchy support (compared to ~10% initial cost for single-level support) seems quite high. Especially right before a new llvm release cut (which is due in 5 days).

I can see how this is needed for implementing the full protocol though (we've already cut on the multi-level callees support), so would you be so kind to hide this behind an option to symbol collector (CollectRefsInMainFile) and then introduce a flag to clangd for controlling it --collect-refs-in-main-file and plumb it to BackgroundIndex and FileIndex via ClangdServer?

[nridge]

In D83536#2143436, @kadircet wrote:

I can see how this is needed for implementing the full protocol though (we've already cut on the multi-level callees support), so would you be so kind to hide this behind an option to symbol collector (CollectRefsInMainFile) and then introduce a flag to clangd for controlling it --collect-refs-in-main-file and plumb it to BackgroundIndex and FileIndex via ClangdServer?

I can do that.

Another thing we could consider, if the space increase is a concern, is to limit which references we store, e.g. to functions only (not variables or classes).

[kadircet]

In D83536#2144282, @nridge wrote:

I can do that.

Another thing we could consider, if the space increase is a concern, is to limit which references we store, e.g. to functions only (not variables or classes).

Yes, that's an option too, but I think putting it behind a flag is the safest for now, as we are cutting close to 11 release.

[nridge]

Rebase

[nridge]

Still need to add the flag.

[nridge]

Add the requested flag

[kadircet]

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

131

please set the default to false

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

219

we don't need the option here right? as IsMainFileOnly flag should always be false for symbols inside headers(preamble).

422

i think we should rather put this option into FileIndex as a field (similar to UseDex) that way we could get rid of some plumbing

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

118

please drop the default

[nridge]

Address review comments

[kadircet]

thanks, lgtm. mostly comments arounds tests, sorry for not looking at them before.

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

156

again default to false.

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

195

this requires it is own separate test, as this is not only testing for indexing two files case now.

I would suggest a minimal test with one header and one cc ref, running background index twice to make sure both configurations work as expected.

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

724–738

could you instead run twice ? once with mainfilerefs = false, and once with it set to true?

815

these seem to be some unrelated formatting changes, please revert before landing.

[nridge]

Address review comments

[nridge]

Rebase

[kadircet]

thanks, LGTM!

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

727

i don't think this is needed, am I missing something? IIRC, inmemoryfs should overwrite any existing files and even if it doesn't we are not really changing file contents in here.

[nridge]

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

727

I initially didn't have this, and got the following assertion in the second runSymbolCollector() call:

ClangdTests: llvm/src/llvm/lib/Support/MemoryBuffer.cpp:48: void llvm::MemoryBuffer::init(const char *, const char *, bool): Assertion `(!RequiresNullTerminator || BufEnd[0] == 0) && "Buffer is not null terminated!"' failed.
 #0 0x00007f4242501c67 llvm::sys::PrintStackTrace(llvm::raw_ostream&) llvm/src/llvm/lib/Support/Unix/Signals.inc:563:11
 #1 0x00007f4242501e09 PrintStackTraceSignalHandler(void*) llvm/src/llvm/lib/Support/Unix/Signals.inc:624:1
 #2 0x00007f424250060b llvm::sys::RunSignalHandlers() llvm/src/llvm/lib/Support/Signals.cpp:67:5
 #3 0x00007f424250251d SignalHandler(int) llvm/src/llvm/lib/Support/Unix/Signals.inc:405:1
 #4 0x00007f42482a10e0 __restore_rt (/lib/x86_64-linux-gnu/libpthread.so.0+0x110e0)
 #5 0x00007f424157dfff raise /build/glibc-77giwP/glibc-2.24/signal/../sysdeps/unix/sysv/linux/raise.c:51:0
 #6 0x00007f424157f42a abort /build/glibc-77giwP/glibc-2.24/stdlib/abort.c:91:0
 #7 0x00007f4241576e67 __assert_fail_base /build/glibc-77giwP/glibc-2.24/assert/assert.c:92:0
 #8 0x00007f4241576f12 (/lib/x86_64-linux-gnu/libc.so.6+0x2bf12)
 #9 0x00007f4242427350 llvm::MemoryBuffer::init(char const*, char const*, bool) llvm/src/llvm/lib/Support/MemoryBuffer.cpp:49:17
#10 0x00007f42424275b4 (anonymous namespace)::MemoryBufferMem<llvm::MemoryBuffer>::MemoryBufferMem(llvm::StringRef, bool) llvm/src/llvm/lib/Support/MemoryBuffer.cpp:89:3
#11 0x00007f4242427400 llvm::MemoryBuffer::getMemBuffer(llvm::StringRef, llvm::StringRef, bool) llvm/src/llvm/lib/Support/MemoryBuffer.cpp:115:7
#12 0x00007f424249ab20 llvm::vfs::detail::(anonymous namespace)::InMemoryFileAdaptor::getBuffer(llvm::Twine const&, long, bool, bool) llvm/src/llvm/lib/Support/VirtualFileSystem.cpp:615:12
#13 0x00007f4242b0ecf1 clang::FileManager::getBufferForFile(clang::FileEntry const*, bool, bool) llvm/src/clang/lib/Basic/FileManager.cpp:474:5
#14 0x00007f4242b52013 clang::SrcMgr::ContentCache::getBuffer(clang::DiagnosticsEngine&, clang::FileManager&, clang::SourceLocation, bool*) const llvm/src/clang/lib/Basic/SourceManager.cpp:172:8
#15 0x00007f4242fe64c0 clang::SourceManager::getBuffer(clang::FileID, clang::SourceLocation, bool*) const llvm/src/clang/include/clang/Basic/SourceManager.h:976:5
#16 0x00007f4242fe2425 clang::Preprocessor::EnterSourceFile(clang::FileID, clang::DirectoryLookup const*, clang::SourceLocation) llvm/src/clang/lib/Lex/PPLexerChange.cpp:77:29
#17 0x00007f424301c104 clang::Preprocessor::EnterMainSourceFile() llvm/src/clang/lib/Lex/Preprocessor.cpp:547:5
#18 0x00007f423d4a23fe clang::ParseAST(clang::Sema&, bool, bool) llvm/src/clang/lib/Parse/ParseAST.cpp:144:33
#19 0x00007f42466df752 clang::ASTFrontendAction::ExecuteAction() llvm/src/clang/lib/Frontend/FrontendAction.cpp:1059:1
#20 0x00007f42466df118 clang::FrontendAction::Execute() llvm/src/clang/lib/Frontend/FrontendAction.cpp:954:7
#21 0x00007f4246655ae4 clang::CompilerInstance::ExecuteAction(clang::FrontendAction&) llvm/src/clang/lib/Frontend/CompilerInstance.cpp:984:23
#22 0x00007f42472a622a clang::tooling::FrontendActionFactory::runInvocation(std::shared_ptr<clang::CompilerInvocation>, clang::FileManager*, std::shared_ptr<clang::PCHContainerOperations>, clang::DiagnosticConsumer*) llvm/src/clang/lib/Tooling/Tooling.cpp:410:14
#23 0x00007f42472a60d4 clang::tooling::ToolInvocation::runInvocation(char const*, clang::driver::Compilation*, std::shared_ptr<clang::CompilerInvocation>, std::shared_ptr<clang::PCHContainerOperations>) llvm/src/clang/lib/Tooling/Tooling.cpp:385:3
#24 0x00007f42472a5015 clang::tooling::ToolInvocation::run() llvm/src/clang/lib/Tooling/Tooling.cpp:370:3
#25 0x0000000000bc916b clang::clangd::(anonymous namespace)::SymbolCollectorTest::runSymbolCollector(llvm::StringRef, llvm::StringRef, std::vector<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, std::allocator<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > > const&) llvm/src/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp:283:15
#26 0x0000000000bda339 clang::clangd::(anonymous namespace)::SymbolCollectorTest_Refs_Test::TestBody() llvm/src/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp:729:3

[kadircet]

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

727

ah makes sense. looks like we are calling MemoryBuffer:getMemBuffer, which would make the buffers unusable after runSymbolCollector exits. I must've been lucky enough to not overwrite the stack.

could you put a comment mentioning the situation? then feel free to land.

clang-tools-extra/clangd/ClangdServer.h

struct Options {
     /// opened files and uses the index to augment code completion results.
     bool BuildDynamicSymbolIndex = false;
     /// Use a heavier and faster in-memory index implementation.
     bool HeavyweightDynamicSymbolIndex = true;
     /// If true, ClangdServer automatically indexes files in the current project
     /// 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;
 
     /// If set, queried to obtain the configuration to handle each request.
     config::Provider *ConfigProvider = nullptr;
 
     /// If set, enable clang-tidy in clangd and use to it get clang-tidy
     /// configurations for a particular file.

clang-tools-extra/clangd/ClangdServer.cpp

ClangdServer::Options::operator TUScheduler::Options() const {
   return Opts;
 }
 
 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            const ThreadsafeFS &TFS, const Options &Opts,
                            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),
       BuildRecoveryAST(Opts.BuildRecoveryAST),
       PreserveRecoveryASTType(Opts.PreserveRecoveryASTType),
       TweakFilter(Opts.TweakFilter), WorkspaceRoot(Opts.WorkspaceRoot),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST

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

 };
 
 // Builds an in-memory index by by running the static indexer action over
 // all commands in a compilation database. Indexing happens in the background.
 // FIXME: it should watch for changes to files on disk.
 class BackgroundIndex : public SwapIndex {
 public:
   struct Options {
     // Arbitrary value to ensure some concurrency in tests.
     // In production an explicit value is specified.
     size_t ThreadPoolSize = 4;
     // Callback that provides notifications as indexing makes progress.
     std::function<void(BackgroundQueue::Stats)> OnProgress = nullptr;
     // Function called to obtain the Context to use while indexing the specified
     // 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.
   /// The current Context will be propagated to each worker thread.
   BackgroundIndex(const ThreadsafeFS &, const GlobalCompilationDatabase &CDB,
                   BackgroundIndexStorage::Factory IndexStorageFactory,
                   Options Opts);
   ~BackgroundIndex(); // Blocks while the current task finishes.
private:
   void update(llvm::StringRef MainFile, IndexFileIn Index,
               const llvm::StringMap<ShardVersion> &ShardVersionsSnapshot,
               bool HadErrors);
 
   // configuration
   const ThreadsafeFS &TFS;
   const GlobalCompilationDatabase &CDB;
   std::function<Context(PathRef)> ContextProvider;
+  bool CollectMainFileRefs;
 
   llvm::Error index(tooling::CompileCommand);
 
   FileSymbols IndexedSymbols;
   BackgroundIndexRebuilder Rebuilder;
   llvm::StringMap<ShardVersion> ShardVersions; // Key is absolute file path.
   std::mutex ShardVersionsMu;
 

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

 
 } // namespace
 
 BackgroundIndex::BackgroundIndex(
     const ThreadsafeFS &TFS, const GlobalCompilationDatabase &CDB,
     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)),
       CommandsChanged(
           CDB.watch([&](const std::vector<std::string> &ChangedFiles) {
             enqueue(ChangedFiles);
           })) {
   assert(Opts.ThreadPoolSize > 0 && "Thread pool size can't be zero.");
IndexOpts.FileFilter = [&ShardVersionsSnapshot](const SourceManager &SM,
     if (!Digest)
       return false;
     auto D = ShardVersionsSnapshot.find(*AbsPath);
     if (D != ShardVersionsSnapshot.end() && D->second.Digest == Digest &&
         !D->second.HadErrors)
       return false; // Skip files that haven't changed, without errors.
     return true;
   };
+  IndexOpts.CollectMainFileRefs = CollectMainFileRefs;
 
   IndexFileIn Index;
   auto Action = createStaticIndexingAction(
       IndexOpts, [&](SymbolSlab S) { Index.Symbols = std::move(S); },
       [&](RefSlab R) { Index.Refs = std::move(R); },
       [&](RelationSlab R) { Index.Relations = std::move(R); },
       [&](IncludeGraph IG) { Index.Sources = std::move(IG); });
 

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

private:
   llvm::StringMap<RefSlabAndCountReferences> RefsSnapshot;
   llvm::StringMap<std::shared_ptr<RelationSlab>> RelatiosSnapshot;
 };
 
 /// This manages symbols from files and an in-memory index on all symbols.
 /// 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.
   void updatePreamble(PathRef Path, llvm::StringRef Version, ASTContext &AST,
                       std::shared_ptr<Preprocessor> PP,
                       const CanonicalIncludes &Includes);
 
   /// Update symbols and references from main file \p Path with
   /// `indexMainDecls`.
   void updateMain(PathRef Path, ParsedAST &AST);
 
 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
   // memory usage.
   // Missing information:
   //  - symbol refs (these are always "from the main file")
   //  - definition locations in the main file
   //
private:
   unsigned PreambleIndexVersion = 0;
 };
 
 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.
 SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
                              std::shared_ptr<Preprocessor> PP,
                              const CanonicalIncludes &Includes);
 
 /// Takes slabs coming from a TU (multiple files) and shards them per

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

 namespace clang {
 namespace clangd {
 namespace {
 
 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.
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
   if (IsIndexMainAST) {
     // We only collect refs when indexing main AST.
FileShardedIndex::getShard(llvm::StringRef Uri) const {
   for (const auto *Rel : It->getValue().Relations) {
     RelB.insert(*Rel);
   }
   IF.Relations = std::move(RelB).build();
   // Explicit move here is needed by some compilers.
   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,
                              std::shared_ptr<Preprocessor> PP,
                              const CanonicalIncludes &Includes) {
   std::vector<Decl *> DeclsToIndex(
       AST.getTranslationUnitDecl()->decls().begin(),
       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,
                          std::unique_ptr<SymbolSlab> Symbols,
                          std::unique_ptr<RefSlab> Refs,
                          std::unique_ptr<RelationSlab> Relations,
                          bool CountReferences) {
   std::lock_guard<std::mutex> Lock(Mutex);
return std::make_unique<dex::Dex>(
         std::move(AllRelations),
         std::make_tuple(std::move(SymbolSlabs), std::move(RefSlabs),
                         std::move(RefsStorage), std::move(SymsStorage)),
         StorageSize);
   }
   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>()) {}
 
 void FileIndex::updatePreamble(PathRef Path, llvm::StringRef Version,
                                ASTContext &AST,
                                std::shared_ptr<Preprocessor> PP,
                                const CanonicalIncludes &Includes) {
   IndexFileIn IF;
auto NewIndex =
     vlog(
         "Build dynamic index for header symbols with estimated memory usage of "
         "{0} bytes",
         PreambleIndex.estimateMemoryUsage());
   }
 }
 
 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))),
       std::make_unique<RelationSlab>(std::move(std::get<2>(Contents))),
       /*CountReferences=*/true);
   size_t IndexVersion = 0;
   auto NewIndex = MainFileSymbols.buildIndex(
       IndexType::Light, DuplicateHandling::Merge, &IndexVersion);
   {

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

struct Options {
     /// Collect macros.
     /// Note that SymbolCollector must be run with preprocessor in order to
     /// collect macros. For example, `indexTopLevelDecls` will not index any
     /// macro even if this is true.
     bool CollectMacro = false;
     /// 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.
     bool StoreAllDocumentation = false;
     /// If this is set, only collect symbols/references from a file if
     /// `FileFilter(SM, FID)` is true. If not set, all files are indexed.
     std::function<bool(const SourceManager &, FileID)> FileFilter = nullptr;
   };

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

bool SymbolCollector::handleDeclOccurrence(
   bool CollectRef = static_cast<bool>(Opts.RefFilter & toRefKind(Roles));
   bool IsOnlyRef =
       !(Roles & (static_cast<unsigned>(index::SymbolRole::Declaration) |
                  static_cast<unsigned>(index::SymbolRole::Definition)));
 
   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()))
     DeclRefs[ND].emplace_back(SM.getFileLoc(Loc), Roles);
   // Don't continue indexing if this is a mere reference.
   if (IsOnlyRef)
     return true;
 

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

desc(
         "User config is from clangd/config.yaml in the following directories:\n"
         "\tWindows: %USERPROFILE%\\AppData\\Local\n"
         "\tMac OS: ~/Library/Preferences/\n"
         "\tOthers: $XDG_CONFIG_HOME, usually ~/.config\n"
         "Configuration is documented at https://clangd.llvm.org/config.html"),
     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",
     cat(Features),
     desc("Address of the remote index server"),
 };
 
 // FIXME(kirillbobyrev): Should this be the location of compile_commands.json?
case PCHStorageFlag::Memory:
     break;
   case PCHStorageFlag::Disk:
     Opts.StorePreamblesInMemory = false;
     break;
   }
   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()) {
     // Load the index asynchronously. Meanwhile SwapIndex returns no results.
     SwapIndex *Placeholder;
     StaticIdx.reset(Placeholder = new SwapIndex(std::make_unique<MemIndex>()));
     AsyncIndexLoad = runAsync<void>([Placeholder] {
       if (auto Idx = loadIndex(IndexFile, /*UseDex=*/true))

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

TEST_F(BackgroundIndexTest, IndexTwoFiles) {
   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++", "-DA=1", testPath("root/A.cc")};
   CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
 
   ASSERT_TRUE(Idx.blockUntilIdleForTest());
   EXPECT_THAT(runFuzzyFind(Idx, ""),
               UnorderedElementsAre(AllOf(Named("common"), NumReferences(1U)),
                                    AllOf(Named("A_CC"), NumReferences(0U)),
                                    AllOf(Named("g"), NumReferences(0U)),
EXPECT_THAT(getRefs(Idx, Common.ID),
               RefsAre({FileURI("unittest:///root/A.h"),
                        FileURI("unittest:///root/A.cc"),
                        FileURI("unittest:///root/B.cc"),
                        FileURI("unittest:///root/B.cc"),
                        FileURI("unittest:///root/B.cc"),
                        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(
       void common();
       void f_b();
       class A_CC {};
       )cpp";
   std::string A_CC = "";

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

EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "Foo").ID,
                                   HaveRanges(Main.ranges("foo")))));
   EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "Bar").ID,
                                   HaveRanges(Main.ranges("bar")))));
   EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "func").ID,
                                   HaveRanges(Main.ranges("func")))));
   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 =
       TestTU::withHeaderCode(SymbolsOnlyInMainCode.code()).headerSymbols();
   EXPECT_THAT(Refs, Contains(Pair(findSymbol(MainSymbols, "a").ID, _)));
   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) {
   Annotations Header(R"(
   #define $foo[[FOO]](X) (X + 1)
   #define $bar[[BAR]](X) (X + 2)
 
   // Macro defined multiple times.

 
 TEST_F(SymbolCollectorTest, SpelledReferences) {
   struct {
     llvm::StringRef Header;
     llvm::StringRef Main;
     llvm::StringRef TargetSymbolName;
   } TestCases[] = {
     {
       R"cpp(
         struct Foo;
         #define MACRO Foo
       )cpp",
       R"cpp(
         struct $spelled[[Foo]] {
           $spelled[[Foo]]();
           ~$spelled[[Foo]]();
         };
TEST_F(SymbolCollectorTest, HeaderAsMainFile) {
   CollectorOpts.RefFilter = RefKind::All;
   Annotations Header(R"(
   class $Foo[[Foo]] {};
 
   void $Func[[Func]]() {
     $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,
               UnorderedElementsAre(Pair(findSymbol(Symbols, "Foo").ID,
                                         HaveRanges(Header.ranges("Foo"))),
                                    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"});
   EXPECT_THAT(Symbols, UnorderedElementsAre(QName("Foo"), QName("Func")));
   EXPECT_THAT(Refs,
               UnorderedElementsAre(Pair(findSymbol(Symbols, "Foo").ID,
                                         HaveRanges(Header.ranges("Foo"))),
                                    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")));
   EXPECT_THAT(Refs,
               UnorderedElementsAre(Pair(findSymbol(Symbols, "Foo").ID,
                                         HaveRanges(Header.ranges("Foo"))),
                                    Pair(findSymbol(Symbols, "Func").ID,
                                         HaveRanges(Header.ranges("Func")))));