[clangd] Collect the number of files referencing a symbol in the static index.

Summary:
This is an important ranking signal.
It's off for the dynamic index for now. Correspondingly, tell the index
infrastructure only to report declarations for the dynamic index.

Reviewers: ioeric, hokein

Subscribers: klimek, ilya-biryukov, jkorous-apple, cfe-commits

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

9 files changed, 89 insertions, 25 deletions

diff --git a/clang-tools-extra/clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp b/clang-tools-extra/clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp
index b4f7f0baeb8..7d62cf12d3c 100644
--- a/clang-tools-extra/clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp
+++ b/clang-tools-extra/clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp
@@ -99,6 +99,7 @@ public:
     auto CollectorOpts = SymbolCollector::Options();
     CollectorOpts.FallbackDir = AssumedHeaderDir;
     CollectorOpts.CollectIncludePath = true;
+    CollectorOpts.CountReferences = true;
     auto Includes = llvm::make_unique<CanonicalIncludes>();
     addSystemHeadersMapping(Includes.get());
     CollectorOpts.Includes = Includes.get();
diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 4f1f10a44c0..db7a8ac43bd 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -26,12 +26,14 @@ std::unique_ptr<SymbolSlab> indexAST(ASTContext &Ctx,
   // AST at this point, but we also need preprocessor callbacks (e.g.
   // CommentHandler for IWYU pragma) to canonicalize includes.
   CollectorOpts.CollectIncludePath = false;
+  CollectorOpts.CountReferences = false;
 
   auto Collector = std::make_shared<SymbolCollector>(std::move(CollectorOpts));
   Collector->setPreprocessor(std::move(PP));
   index::IndexingOptions IndexOpts;
+  // We only need declarations, because we don't count references.
   IndexOpts.SystemSymbolFilter =
-      index::IndexingOptions::SystemSymbolFilterKind::All;
+      index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
 
   index::indexTopLevelDecls(Ctx, Decls, Collector, IndexOpts);
diff --git a/clang-tools-extra/clangd/index/Index.h b/clang-tools-extra/clangd/index/Index.h
index 629db530928..efd18d9b940 100644
--- a/clang-tools-extra/clangd/index/Index.h
+++ b/clang-tools-extra/clangd/index/Index.h
@@ -131,6 +131,9 @@ struct Symbol {
   //   * For non-inline functions, the canonical declaration typically appears
   //     in the ".h" file corresponding to the definition.
   SymbolLocation CanonicalDeclaration;
+  // The number of translation units that reference this symbol from their main
+  // file. This number is only meaningful if aggregated in an index.
+  unsigned References = 0;
 
   /// A brief description of the symbol that can be displayed in the completion
   /// candidate list. For example, "Foo(X x, Y y) const" is a labal for a
diff --git a/clang-tools-extra/clangd/index/Merge.cpp b/clang-tools-extra/clangd/index/Merge.cpp
index 57c62c5bc5c..9df5376bbfd 100644
--- a/clang-tools-extra/clangd/index/Merge.cpp
+++ b/clang-tools-extra/clangd/index/Merge.cpp
@@ -73,6 +73,7 @@ mergeSymbol(const Symbol &L, const Symbol &R, Symbol::Details *Scratch) {
     S.Definition = O.Definition;
   if (!S.CanonicalDeclaration)
     S.CanonicalDeclaration = O.CanonicalDeclaration;
+  S.References += O.References;
   if (S.CompletionLabel == "")
     S.CompletionLabel = O.CompletionLabel;
   if (S.CompletionFilterText == "")
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.cpp b/clang-tools-extra/clangd/index/SymbolCollector.cpp
index e237e136d25..b27ace5f3ff 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ b/clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -228,37 +228,58 @@ bool SymbolCollector::handleDeclOccurence(
     ArrayRef<index::SymbolRelation> Relations, FileID FID, unsigned Offset,
     index::IndexDataConsumer::ASTNodeInfo ASTNode) {
   assert(ASTCtx && PP.get() && "ASTContext and Preprocessor must be set.");
+  assert(CompletionAllocator && CompletionTUInfo);
+  const NamedDecl *ND = llvm::dyn_cast<NamedDecl>(D);
+  if (!ND)
+    return true;
+
+  // Mark D as referenced if this is a reference coming from the main file.
+  // D may not be an interesting symbol, but it's cheaper to check at the end.
+  if (Opts.CountReferences &&
+      (Roles & static_cast<unsigned>(index::SymbolRole::Reference)) &&
+      ASTCtx->getSourceManager().getMainFileID() == FID)
+    ReferencedDecls.insert(ND);
 
-  // FIXME: collect all symbol references.
+  // Don't continue indexing if this is a mere reference.
   if (!(Roles & static_cast<unsigned>(index::SymbolRole::Declaration) ||
         Roles & static_cast<unsigned>(index::SymbolRole::Definition)))
     return true;
+  if (shouldFilterDecl(ND, ASTCtx, Opts))
+    return true;
 
-  assert(CompletionAllocator && CompletionTUInfo);
+  llvm::SmallString<128> USR;
+  if (index::generateUSRForDecl(ND, USR))
+    return true;
+  SymbolID ID(USR);
 
-  if (const NamedDecl *ND = llvm::dyn_cast<NamedDecl>(D)) {
-    if (shouldFilterDecl(ND, ASTCtx, Opts))
-      return true;
+  const NamedDecl &OriginalDecl = *cast<NamedDecl>(ASTNode.OrigD);
+  const Symbol *BasicSymbol = Symbols.find(ID);
+  if (!BasicSymbol) // Regardless of role, ND is the canonical declaration.
+    BasicSymbol = addDeclaration(*ND, std::move(ID));
+  else if (isPreferredDeclaration(OriginalDecl, Roles))
+    // If OriginalDecl is preferred, replace the existing canonical
+    // declaration (e.g. a class forward declaration). There should be at most
+    // one duplicate as we expect to see only one preferred declaration per
+    // TU, because in practice they are definitions.
+    BasicSymbol = addDeclaration(OriginalDecl, std::move(ID));
+
+  if (Roles & static_cast<unsigned>(index::SymbolRole::Definition))
+    addDefinition(OriginalDecl, *BasicSymbol);
+  return true;
+}
+
+void SymbolCollector::finish() {
+  // At the end of the TU, add 1 to the refcount of the ReferencedDecls.
+  for (const auto *ND : ReferencedDecls) {
     llvm::SmallString<128> USR;
-    if (index::generateUSRForDecl(ND, USR))
-      return true;
-
-    const NamedDecl &OriginalDecl = *cast<NamedDecl>(ASTNode.OrigD);
-    auto ID = SymbolID(USR);
-    const Symbol *BasicSymbol = Symbols.find(ID);
-    if (!BasicSymbol) // Regardless of role, ND is the canonical declaration.
-      BasicSymbol = addDeclaration(*ND, std::move(ID));
-    else if (isPreferredDeclaration(OriginalDecl, Roles))
-      // If OriginalDecl is preferred, replace the existing canonical
-      // declaration (e.g. a class forward declaration). There should be at most
-      // one duplicate as we expect to see only one preferred declaration per
-      // TU, because in practice they are definitions.
-      BasicSymbol = addDeclaration(OriginalDecl, std::move(ID));
-
-    if (Roles & static_cast<unsigned>(index::SymbolRole::Definition))
-      addDefinition(OriginalDecl, *BasicSymbol);
+    if (!index::generateUSRForDecl(ND, USR))
+      if (const auto *S = Symbols.find(SymbolID(USR))) {
+        Symbol Inc = *S;
+        ++Inc.References;
+        Symbols.insert(Inc);
+      }
   }
-  return true;
+  ReferencedDecls.clear();
 }
 
 const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND,
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.h b/clang-tools-extra/clangd/index/SymbolCollector.h
index c18f74a8362..7237e391ba0 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.h
+++ b/clang-tools-extra/clangd/index/SymbolCollector.h
@@ -45,6 +45,8 @@ public:
     /// 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;
   };
 
   SymbolCollector(Options Opts);
@@ -63,6 +65,8 @@ public:
 
   SymbolSlab takeSymbols() { return std::move(Symbols).build(); }
 
+  void finish() override;
+
 private:
   const Symbol *addDeclaration(const NamedDecl &, SymbolID);
   void addDefinition(const NamedDecl &, const Symbol &DeclSymbol);
@@ -74,6 +78,8 @@ private:
   std::shared_ptr<GlobalCodeCompletionAllocator> CompletionAllocator;
   std::unique_ptr<CodeCompletionTUInfo> CompletionTUInfo;
   Options Opts;
+  // Decls referenced from the current TU, flushed on finish().
+  llvm::DenseSet<const NamedDecl *> ReferencedDecls;
 };
 
 } // namespace clangd
diff --git a/clang-tools-extra/clangd/index/SymbolYAML.cpp b/clang-tools-extra/clangd/index/SymbolYAML.cpp
index 379cb334d6b..307b0af5b5c 100644
--- a/clang-tools-extra/clangd/index/SymbolYAML.cpp
+++ b/clang-tools-extra/clangd/index/SymbolYAML.cpp
@@ -100,6 +100,7 @@ template <> struct MappingTraits<Symbol> {
     IO.mapOptional("CanonicalDeclaration", Sym.CanonicalDeclaration,
                    SymbolLocation());
     IO.mapOptional("Definition", Sym.Definition, SymbolLocation());
+    IO.mapOptional("References", Sym.References, 0u);
     IO.mapRequired("CompletionLabel", Sym.CompletionLabel);
     IO.mapRequired("CompletionFilterText", Sym.CompletionFilterText);
     IO.mapRequired("CompletionPlainInsertText", Sym.CompletionPlainInsertText);
diff --git a/clang-tools-extra/unittests/clangd/IndexTests.cpp b/clang-tools-extra/unittests/clangd/IndexTests.cpp
index 20eb4527d9f..5e5ddbfb86e 100644
--- a/clang-tools-extra/unittests/clangd/IndexTests.cpp
+++ b/clang-tools-extra/unittests/clangd/IndexTests.cpp
@@ -225,6 +225,8 @@ TEST(MergeTest, Merge) {
   L.Name = R.Name = "Foo";                    // same in both
   L.CanonicalDeclaration.FileURI = "file:///left.h"; // differs
   R.CanonicalDeclaration.FileURI = "file:///right.h";
+  L.References = 1;
+  R.References = 2;
   L.CompletionPlainInsertText = "f00";        // present in left only
   R.CompletionSnippetInsertText = "f0{$1:0}"; // present in right only
   Symbol::Details DetL, DetR;
@@ -238,6 +240,7 @@ TEST(MergeTest, Merge) {
   Symbol M = mergeSymbol(L, R, &Scratch);
   EXPECT_EQ(M.Name, "Foo");
   EXPECT_EQ(M.CanonicalDeclaration.FileURI, "file:///left.h");
+  EXPECT_EQ(M.References, 3u);
   EXPECT_EQ(M.CompletionPlainInsertText, "f00");
   EXPECT_EQ(M.CompletionSnippetInsertText, "f0{$1:0}");
   ASSERT_TRUE(M.Detail);
diff --git a/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp b/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp
index c6a155a4a36..76eace528f4 100644
--- a/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp
+++ b/clang-tools-extra/unittests/clangd/SymbolCollectorTests.cpp
@@ -59,6 +59,7 @@ MATCHER_P(DefRange, Offsets, "") {
   return arg.Definition.StartOffset == Offsets.first &&
          arg.Definition.EndOffset == Offsets.second;
 }
+MATCHER_P(Refs, R, "") { return int(arg.References) == R; }
 
 namespace clang {
 namespace clangd {
@@ -201,7 +202,7 @@ TEST_F(SymbolCollectorTest, CollectSymbols) {
 TEST_F(SymbolCollectorTest, Template) {
   Annotations Header(R"(
     // Template is indexed, specialization and instantiation is not.
-    template <class T> struct [[Tmpl]] {T x = 0};
+    template <class T> struct [[Tmpl]] {T x = 0;};
     template <> struct Tmpl<int> {};
     extern template struct Tmpl<float>;
     template struct Tmpl<double>;
@@ -242,6 +243,31 @@ TEST_F(SymbolCollectorTest, Locations) {
           AllOf(QName("Z"), DeclRange(Header.offsetRange("zdecl")))));
 }
 
+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
+  )";
+  const std::string Main = R"(
+    W* w = nullptr;
+    W* w2 = nullptr; // only one usage counts
+    X x();
+    class V;
+    V* v = nullptr; // Used, but not eligible for indexing.
+    class Y{}; // definition doesn't count as a reference
+  )";
+  CollectorOpts.CountReferences = true;
+  runSymbolCollector(Header, Main);
+  EXPECT_THAT(Symbols,
+              UnorderedElementsAre(AllOf(QName("W"), Refs(1)),
+                                   AllOf(QName("X"), Refs(1)),
+                                   AllOf(QName("Y"), Refs(0)),
+                                   AllOf(QName("Z"), Refs(0)), QName("y")));
+}
+
 TEST_F(SymbolCollectorTest, SymbolRelativeNoFallback) {
   runSymbolCollector("class Foo {};", /*Main=*/"");
   EXPECT_THAT(Symbols, UnorderedElementsAre(