[clangd] Collect refs from headers.

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).

Reviewers: sammccall

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

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

llvm-svn: 344678

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

@@ -66,8 +66,10 @@ createStaticIndexingAction(SymbolCollector::Options Opts,
   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();

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

@@ -21,9 +21,8 @@ namespace clangd {
 // 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,

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

@@ -19,6 +19,7 @@
 #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"
@@ -354,7 +355,8 @@ bool SymbolCollector::handleDeclOccurence(
     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)
@@ -474,26 +476,43 @@ void SymbolCollector::finish() {
   }
 
   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)) {
+  if (auto MainFileURI = GetURI(SM.getMainFileID())) {
-    std::string MainURI = *MainFileURI;
     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);
-          auto Range =
+          if (auto FileURI = GetURI(FileID)) {
-              getTokenRange(LocAndRole.first, SM, ASTCtx->getLangOpts());
+            auto Range =
-          R.Location.Start = Range.first;
+                getTokenRange(LocAndRole.first, SM, ASTCtx->getLangOpts());
-          R.Location.End = Range.second;
+            Ref R;
-          R.Location.FileURI = MainURI;
+            R.Location.Start = Range.first;
-          R.Kind = toRefKind(LocAndRole.second);
+            R.Location.End = Range.second;
-          Refs.insert(*ID, R);
+            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();

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

@@ -57,6 +57,11 @@ public:
     /// 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.

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

@@ -479,6 +479,17 @@ TEST_F(SymbolCollectorTest, Refs) {
   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;