[clangd] Collect refs from headers.

[hokein]

Summary

Add a flag to SymbolCollector to collect refs fdrom headers.

Note that we collect refs from headers in static index, and we don't do it for
dynamic index because of the preamble (we skip function body in preamble,
collecting it will result incomplete results).

[hokein]

For references, the binary-format index file (for whole llvm project) size:

	Before	WithRef
size	50MB	91MB
memory	255MB	412MB

There are 300,000 symbols.

[sammccall]

This looks reasonable. +80% to index size is a shame, but not disastrous, and we have ways to shave it.

Can you check:

• how much we're increasing the in-memory size (Dex)
• that we're not outputting duplicate refs when preambles overlap between TUs

clangd/index/IndexAction.h

24

just "all references are collected"?

clangd/index/SymbolCollector.cpp

498

maybe pull out

auto GetURI = [&](clang::FileID FID) -> Optional<StringRef> {
  ... look up or populate cache entry ...
}

then both the outer if and inner lookups can be written as:
if (auto MainFileURI = GetURI(SM.getMainFileID())) {
etc

506

I think you probably also want to cache this case - using an empty string is fine

clangd/index/SymbolCollector.h

62

available -> meaningful

63

having "true" mean don't do something is a bit confusing.

What about bool RefsInHeaders = false?

[sammccall]

(please do check there are no duplicates in the output)

[hokein]

Address review comments.

[hokein]

In D53322#1266536, @sammccall wrote:

(please do check there are no duplicates in the output)

We do deduplication when building the RefSlab. And double-checked with the output, no duplications there.

[hokein]

minor cleanup.

clangd/index/IndexAction.h

 namespace clangd {
 
 // Creates an action that indexes translation units and delivers the results
 // for SymbolsCallback (each slab corresponds to one TU).
 //
 // Only a subset of SymbolCollector::Options are respected:
 //   - include paths are always collected, and canonicalized appropriately
 //   - references are always counted
-//   - main-file refs are collected (if RefsCallback is non-null)
+//   - all references are collected (if RefsCallback is non-null)
 //   - the symbol origin is always Static
-// FIXME: refs from headers should also be collected.
 std::unique_ptr<FrontendAction>
 createStaticIndexingAction(SymbolCollector::Options Opts,
                            std::function<void(SymbolSlab)> SymbolsCallback,
                            std::function<void(RefSlab)> RefsCallback);
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/index/IndexAction.cpp

createStaticIndexingAction(SymbolCollector::Options Opts,
                            std::function<void(SymbolSlab)> SymbolsCallback,
                            std::function<void(RefSlab)> RefsCallback) {
   index::IndexingOptions IndexOpts;
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   Opts.CollectIncludePath = true;
   Opts.CountReferences = true;
   Opts.Origin = SymbolOrigin::Static;
-  if (RefsCallback != nullptr)
+  if (RefsCallback != nullptr) {
     Opts.RefFilter = RefKind::All;
+    Opts.RefsInHeaders = true;
+  }
   auto Includes = llvm::make_unique<CanonicalIncludes>();
   addSystemHeadersMapping(Includes.get());
   Opts.Includes = Includes.get();
   return llvm::make_unique<IndexAction>(
       std::make_shared<SymbolCollector>(std::move(Opts)), std::move(Includes),
       IndexOpts, SymbolsCallback, RefsCallback);
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/SymbolCollector.h

struct Options {
     /// If set, this is used to map symbol #include path to a potentially
     /// different #include path.
     const CanonicalIncludes *Includes = nullptr;
     // Populate the Symbol.References field.
     bool CountReferences = false;
     /// The symbol ref kinds that will be collected.
     /// If not set, SymbolCollector will not collect refs.
     RefKind RefFilter = RefKind::Unknown;
+    /// If set to true, SymbolCollector will collect all refs (from main file
+    /// and included headers); otherwise, only refs from main file will be
+    /// collected.
+    /// This flag is only meaningful when RefFilter is set.
+    bool RefsInHeaders = false;
     // Every symbol collected will be stamped with this origin.
     SymbolOrigin Origin = SymbolOrigin::Unknown;
     /// 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;
   };

clangd/index/SymbolCollector.cpp

 #include "CodeCompletionStrings.h"
 #include "Logger.h"
 #include "SourceCode.h"
 #include "URI.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclTemplate.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Index/USRGeneration.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
bool SymbolCollector::handleDeclOccurence(
   bool IsOnlyRef =
       !(Roles & (static_cast<unsigned>(index::SymbolRole::Declaration) |
                  static_cast<unsigned>(index::SymbolRole::Definition)));
 
   if (IsOnlyRef && !CollectRef)
     return true;
   if (!shouldCollectSymbol(*ND, *ASTCtx, Opts))
     return true;
-  if (CollectRef && SM.getFileID(SpellingLoc) == SM.getMainFileID())
+  if (CollectRef &&
+      (Opts.RefsInHeaders || SM.getFileID(SpellingLoc) == SM.getMainFileID()))
     DeclRefs[ND].emplace_back(SpellingLoc, Roles);
   // Don't continue indexing if this is a mere reference.
   if (IsOnlyRef)
     return true;
 
   auto ID = getSymbolID(ND);
   if (!ID)
     return true;
if (Opts.CollectMacro) {
     for (const IdentifierInfo *II : ReferencedMacros) {
       if (const auto *MI = PP->getMacroDefinition(II).getMacroInfo())
         if (auto ID = getSymbolID(*II, MI, PP->getSourceManager()))
           IncRef(*ID);
     }
   }
 
   const auto &SM = ASTCtx->getSourceManager();
-  auto* MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
+  llvm::DenseMap<FileID, std::string> URICache;
+  auto GetURI = [&](FileID FID) -> llvm::Optional<std::string> {
+    auto Found = URICache.find(FID);
+    if (Found == URICache.end()) {
+      // Ignore cases where we can not find a corresponding file entry
+      // for the loc, thoses are not interesting, e.g. symbols formed
+      // via macro concatenation.
+      if (auto *FileEntry = SM.getFileEntryForID(FID)) {
+        auto FileURI = toURI(SM, FileEntry->getName(), Opts);
+        if (!FileURI) {
+          log("Failed to create URI for file: {0}\n", FileEntry);
+          FileURI = ""; // reset to empty as we also want to cache this case.
+        }
+        Found = URICache.insert({FID, *FileURI}).first;
+      }
+    }
+    return Found->second;
+  };
 
-  if (auto MainFileURI = toURI(SM, MainFileEntry->getName(), Opts)) {
-    std::string MainURI = *MainFileURI;
+  if (auto MainFileURI = GetURI(SM.getMainFileID())) {
     for (const auto &It : DeclRefs) {
       if (auto ID = getSymbolID(It.first)) {
         for (const auto &LocAndRole : It.second) {
-          Ref R;
+          auto FileID = SM.getFileID(LocAndRole.first);
+          if (auto FileURI = GetURI(FileID)) {
             auto Range =
                 getTokenRange(LocAndRole.first, SM, ASTCtx->getLangOpts());
+            Ref R;
             R.Location.Start = Range.first;
             R.Location.End = Range.second;
-          R.Location.FileURI = MainURI;
+            R.Location.FileURI = *FileURI;
             R.Kind = toRefKind(LocAndRole.second);
             Refs.insert(*ID, R);
           }
         }
       }
-  } else {
-    log("Failed to create URI for main file: {0}", MainFileEntry->getName());
+    }
   }
 
   ReferencedDecls.clear();
   ReferencedMacros.clear();
   DeclRefs.clear();
 }
 
 const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND,

unittests/clangd/SymbolCollectorTests.cpp

TEST_F(SymbolCollectorTest, Refs) {
   // Symbols *only* in the main file (a, b, c) had no refs collected.
   auto MainSymbols =
       TestTU::withHeaderCode(SymbolsOnlyInMainCode.code()).headerSymbols();
   EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(MainSymbols, "a").ID, _))));
   EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(MainSymbols, "b").ID, _))));
   EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(MainSymbols, "c").ID, _))));
 }
 
+TEST_F(SymbolCollectorTest, RefsInHeaders) {
+  CollectorOpts.RefFilter = RefKind::All;
+  CollectorOpts.RefsInHeaders = true;
+  Annotations Header(R"(
+  class [[Foo]] {};
+  )");
+  runSymbolCollector(Header.code(), "");
+  EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "Foo").ID,
+                                  HaveRanges(Header.ranges()))));
+}
+
 TEST_F(SymbolCollectorTest, References) {
   const std::string Header = R"(
     class W;
     class X {};
     class Y;
     class Z {}; // not used anywhere
     Y* y = nullptr;  // used in header doesn't count
     #define GLOBAL_Z(name) Z name;