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

[sammccall]

Summary

This is an important ranking signal.
It's off for the dynamic index for now.

clangd/global-symbol-builder/GlobalSymbolBuilderMain.cpp

clang::FrontendAction *create() override {
 
     index::IndexingOptions IndexOpts;
     IndexOpts.SystemSymbolFilter =
         index::IndexingOptions::SystemSymbolFilterKind::All;
     IndexOpts.IndexFunctionLocals = false;
     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();
     return new WrappedIndexAction(
         std::make_shared<SymbolCollector>(std::move(CollectorOpts)),
         std::move(Includes), IndexOpts, Ctx);
   }
 

clangd/index/Index.h

struct Symbol {
   // This just covers the symbol name.
   // This may be the same as Definition.
   //
   // A C++ symbol may have multiple declarations, and we pick one to prefer.
   //   * For classes, the canonical declaration should be the definition.
   //   * 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
   /// function.
   llvm::StringRef CompletionLabel;
   /// The piece of text that the user is expected to type to match the
   /// code-completion string, typically a keyword or the name of a declarator or
   /// macro.

clangd/index/Merge.cpp

mergeSymbol(const Symbol &L, const Symbol &R, Symbol::Details *Scratch) {
   const Symbol &O = PreferR ? L : R; // The "other" less-preferred symbol.
 
   // For each optional field, fill it from O if missing in S.
   // (It might be missing in O too, but that's a no-op).
   if (!S.Definition)
     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 == "")
     S.CompletionFilterText = O.CompletionFilterText;
   if (S.CompletionPlainInsertText == "")
     S.CompletionPlainInsertText = O.CompletionPlainInsertText;
   if (S.CompletionSnippetInsertText == "")
     S.CompletionSnippetInsertText = O.CompletionSnippetInsertText;

clangd/index/SymbolCollector.h

struct Options {
     /// Specifies URI schemes that can be used to generate URIs for file paths
     /// in symbols. The list of schemes will be tried in order until a working
     /// scheme is found. If no scheme works, symbol location will be dropped.
     std::vector<std::string> URISchemes = {"file"};
     bool CollectIncludePath = false;
     /// 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);
 
   void initialize(ASTContext &Ctx) override;
 
   void setPreprocessor(std::shared_ptr<Preprocessor> PP) override {
     this->PP = std::move(PP);
   }
 
   bool
   handleDeclOccurence(const Decl *D, index::SymbolRoleSet Roles,
                       ArrayRef<index::SymbolRelation> Relations, FileID FID,
                       unsigned Offset,
                       index::IndexDataConsumer::ASTNodeInfo ASTNode) override;
 
   SymbolSlab takeSymbols() { return std::move(Symbols).build(); }
 
+  void finish() override;
+
 private:
   const Symbol *addDeclaration(const NamedDecl &, SymbolID);
   void addDefinition(const NamedDecl &, const Symbol &DeclSymbol);
 
   // All Symbols collected from the AST.
   SymbolSlab::Builder Symbols;
   ASTContext *ASTCtx;
   std::shared_ptr<Preprocessor> PP;
   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
 } // namespace clang

clangd/index/SymbolCollector.cpp

 }
 
 // Always return true to continue indexing.
 bool SymbolCollector::handleDeclOccurence(
     const Decl *D, index::SymbolRoleSet Roles,
     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;
 
-  // FIXME: collect all symbol references.
+  // 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);
+
+  // 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;
-
-  assert(CompletionAllocator && CompletionTUInfo);
-
-  if (const NamedDecl *ND = llvm::dyn_cast<NamedDecl>(D)) {
-    if (shouldFilterDecl(ND, ASTCtx, Opts))
-      return true;
-    llvm::SmallString<128> USR;
-    if (index::generateUSRForDecl(ND, USR))
-      return true;
+
+  SymbolID ID(USR);
 
-    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);
-  }
   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))
+      if (const auto *S = Symbols.find(SymbolID(USR))) {
+        Symbol Inc = *S;
+        ++Inc.References;
+        Symbols.insert(Inc);
+      }
+  }
+  ReferencedDecls.clear();
+}
+
 const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND,
                                               SymbolID ID) {
   auto &SM = ND.getASTContext().getSourceManager();
 
   std::string QName;
   llvm::raw_string_ostream OS(QName);
   PrintingPolicy Policy(ASTCtx->getLangOpts());
   // Note that inline namespaces are treated as transparent scopes. This

clangd/index/SymbolYAML.cpp

MappingNormalization<ArenaPtr<Symbol::Details>, const Symbol::Details *>
         NDetail(IO, Sym.Detail);
     IO.mapRequired("ID", NSymbolID->HexString);
     IO.mapRequired("Name", Sym.Name);
     IO.mapRequired("Scope", Sym.Scope);
     IO.mapRequired("SymInfo", Sym.SymInfo);
     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);
 
     IO.mapOptional("CompletionSnippetInsertText",
                    Sym.CompletionSnippetInsertText);
     IO.mapOptional("Detail", NDetail->Opt);
   }

unittests/clangd/IndexTests.cpp

 }
 
 TEST(MergeTest, Merge) {
   Symbol L, R;
   L.ID = R.ID = SymbolID("hello");
   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;
   DetL.CompletionDetail = "DetL";
   DetR.CompletionDetail = "DetR";
   DetR.Documentation = "--doc--";
   L.Detail = &DetL;
   R.Detail = &DetR;
 
   Symbol::Details Scratch;
   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);
   EXPECT_EQ(M.Detail->CompletionDetail, "DetL");
   EXPECT_EQ(M.Detail->Documentation, "--doc--");
 }
 
 TEST(MergeTest, PreferSymbolWithDefn) {

unittests/clangd/SymbolCollectorTests.cpp

 MATCHER_P(DeclRange, Offsets, "") {
   return arg.CanonicalDeclaration.StartOffset == Offsets.first &&
       arg.CanonicalDeclaration.EndOffset == Offsets.second;
 }
 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 {
 
 namespace {
 class SymbolIndexActionFactory : public tooling::FrontendActionFactory {
 public:
   SymbolIndexActionFactory(SymbolCollector::Options COpts,
EXPECT_THAT(Symbols,
                    QName("kStr"), QName("foo"), QName("foo::bar"),
                    QName("foo::int32"), QName("foo::int32_t"), QName("foo::v1"),
                    QName("foo::bar::v2"), QName("foo::baz")}));
 }
 
 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>;
   )");
   runSymbolCollector(Header.code(), /*Main=*/"");
   EXPECT_THAT(Symbols, UnorderedElementsAreArray({AllOf(
                            QName("Tmpl"), DeclRange(Header.offsetRange()))}));
 }
EXPECT_THAT(
                 DefRange(Main.offsetRange("xdef"))),
           AllOf(QName("Cls"), DeclRange(Header.offsetRange("clsdecl")),
                 DefRange(Main.offsetRange("clsdef"))),
           AllOf(QName("print"), DeclRange(Header.offsetRange("printdecl")),
                 DefRange(Main.offsetRange("printdef"))),
           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(
                            AllOf(QName("Foo"), DeclURI(TestHeaderURI))));
 }
 
 TEST_F(SymbolCollectorTest, SymbolRelativeWithFallback) {
   TestHeaderName = "x.h";