[clangd] Count number of references while merging RefSlabs inside FileIndex

Summary:
For counting number of references clangd was relying on merging every
duplication of a symbol. Unfortunately this does not apply to FileIndex(and one
of its users' BackgroundIndex), since we get rid of duplication by simply
dropping symbols coming from non-canonical locations. So only one or two(coming
from canonical declaration header and defined source file, if exists)
replications of the same symbol reaches merging step.

This patch changes reference counting logic to rather count number of different
RefSlabs a given SymbolID exists.

Reviewers: ilya-biryukov

Subscribers: ioeric, MaskRay, jkorous, mgrang, arphaman, jdoerfert, cfe-commits

Tags: #clang

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

llvm-svn: 360344

7 files changed, 104 insertions, 33 deletions

diff --git a/clang-tools-extra/clangd/index/Background.cpp b/clang-tools-extra/clangd/index/Background.cpp
index 67053948a03..02cd11d1389 100644
--- a/clang-tools-extra/clangd/index/Background.cpp
+++ b/clang-tools-extra/clangd/index/Background.cpp
@@ -356,7 +356,8 @@ void BackgroundIndex::update(llvm::StringRef MainFile, IndexFileIn Index,
       // This can override a newer version that is added in another thread, if
       // this thread sees the older version but finishes later. This should be
       // rare in practice.
-      IndexedSymbols.update(Path, std::move(SS), std::move(RS));
+      IndexedSymbols.update(Path, std::move(SS), std::move(RS),
+                            Path == MainFile);
     }
   }
 }
@@ -478,7 +479,8 @@ BackgroundIndex::loadShard(const tooling::CompileCommand &Cmd,
   struct ShardInfo {
     std::string AbsolutePath;
     std::unique_ptr<IndexFileIn> Shard;
-    FileDigest Digest;
+    FileDigest Digest = {};
+    bool CountReferences = false;
   };
   std::vector<ShardInfo> IntermediateSymbols;
   // Make sure we don't have duplicate elements in the queue. Keys are absolute
@@ -539,6 +541,7 @@ BackgroundIndex::loadShard(const tooling::CompileCommand &Cmd,
       SI.AbsolutePath = CurDependency.Path;
       SI.Shard = std::move(Shard);
       SI.Digest = I.getValue().Digest;
+      SI.CountReferences = I.getValue().IsTU;
       IntermediateSymbols.push_back(std::move(SI));
       // Check if the source needs re-indexing.
       // Get the digest, skip it if file doesn't exist.
@@ -568,7 +571,8 @@ BackgroundIndex::loadShard(const tooling::CompileCommand &Cmd,
                     ? llvm::make_unique<RefSlab>(std::move(*SI.Shard->Refs))
                     : nullptr;
       IndexedFileDigests[SI.AbsolutePath] = SI.Digest;
-      IndexedSymbols.update(SI.AbsolutePath, std::move(SS), std::move(RS));
+      IndexedSymbols.update(SI.AbsolutePath, std::move(SS), std::move(RS),
+                            SI.CountReferences);
     }
   }
 
diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 7eca85e22e0..22c6e860fb8 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -90,28 +90,36 @@ SymbolSlab indexHeaderSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
 }
 
 void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Symbols,
-                         std::unique_ptr<RefSlab> Refs) {
+                         std::unique_ptr<RefSlab> Refs, bool CountReferences) {
   std::lock_guard<std::mutex> Lock(Mutex);
   if (!Symbols)
     FileToSymbols.erase(Path);
   else
     FileToSymbols[Path] = std::move(Symbols);
-  if (!Refs)
+  if (!Refs) {
     FileToRefs.erase(Path);
-  else
-    FileToRefs[Path] = std::move(Refs);
+    return;
+  }
+  RefSlabAndCountReferences Item;
+  Item.CountReferences = CountReferences;
+  Item.Slab = std::move(Refs);
+  FileToRefs[Path] = std::move(Item);
 }
 
 std::unique_ptr<SymbolIndex>
 FileSymbols::buildIndex(IndexType Type, DuplicateHandling DuplicateHandle) {
   std::vector<std::shared_ptr<SymbolSlab>> SymbolSlabs;
   std::vector<std::shared_ptr<RefSlab>> RefSlabs;
+  std::vector<RefSlab *> MainFileRefs;
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     for (const auto &FileAndSymbols : FileToSymbols)
       SymbolSlabs.push_back(FileAndSymbols.second);
-    for (const auto &FileAndRefs : FileToRefs)
-      RefSlabs.push_back(FileAndRefs.second);
+    for (const auto &FileAndRefs : FileToRefs) {
+      RefSlabs.push_back(FileAndRefs.second.Slab);
+      if (FileAndRefs.second.CountReferences)
+        MainFileRefs.push_back(RefSlabs.back().get());
+    }
   }
   std::vector<const Symbol *> AllSymbols;
   std::vector<Symbol> SymsStorage;
@@ -120,25 +128,35 @@ FileSymbols::buildIndex(IndexType Type, DuplicateHandling DuplicateHandle) {
     llvm::DenseMap<SymbolID, Symbol> Merged;
     for (const auto &Slab : SymbolSlabs) {
       for (const auto &Sym : *Slab) {
+        assert(Sym.References == 0 &&
+               "Symbol with non-zero references sent to FileSymbols");
         auto I = Merged.try_emplace(Sym.ID, Sym);
         if (!I.second)
           I.first->second = mergeSymbol(I.first->second, Sym);
       }
     }
+    for (const RefSlab *Refs : MainFileRefs)
+      for (const auto &Sym : *Refs) {
+        auto It = Merged.find(Sym.first);
+        assert(It != Merged.end() && "Reference to unknown symbol");
+        It->getSecond().References += Sym.second.size();
+      }
     SymsStorage.reserve(Merged.size());
     for (auto &Sym : Merged) {
       SymsStorage.push_back(std::move(Sym.second));
       AllSymbols.push_back(&SymsStorage.back());
     }
-    // FIXME: aggregate symbol reference count based on references.
     break;
   }
   case DuplicateHandling::PickOne: {
     llvm::DenseSet<SymbolID> AddedSymbols;
     for (const auto &Slab : SymbolSlabs)
-      for (const auto &Sym : *Slab)
+      for (const auto &Sym : *Slab) {
+        assert(Sym.References == 0 &&
+               "Symbol with non-zero references sent to FileSymbols");
         if (AddedSymbols.insert(Sym.ID).second)
           AllSymbols.push_back(&Sym);
+      }
     break;
   }
   }
@@ -201,9 +219,9 @@ void FileIndex::updatePreamble(PathRef Path, ASTContext &AST,
                                std::shared_ptr<Preprocessor> PP,
                                const CanonicalIncludes &Includes) {
   auto Symbols = indexHeaderSymbols(AST, std::move(PP), Includes);
-  PreambleSymbols.update(Path,
-                         llvm::make_unique<SymbolSlab>(std::move(Symbols)),
-                         llvm::make_unique<RefSlab>());
+  PreambleSymbols.update(
+      Path, llvm::make_unique<SymbolSlab>(std::move(Symbols)),
+      llvm::make_unique<RefSlab>(), /*CountReferences=*/false);
   PreambleIndex.reset(
       PreambleSymbols.buildIndex(UseDex ? IndexType::Heavy : IndexType::Light,
                                  DuplicateHandling::PickOne));
@@ -213,7 +231,8 @@ void FileIndex::updateMain(PathRef Path, ParsedAST &AST) {
   auto Contents = indexMainDecls(AST);
   MainFileSymbols.update(
       Path, llvm::make_unique<SymbolSlab>(std::move(Contents.first)),
-      llvm::make_unique<RefSlab>(std::move(Contents.second)));
+      llvm::make_unique<RefSlab>(std::move(Contents.second)),
+      /*CountReferences=*/true);
   MainFileIndex.reset(
       MainFileSymbols.buildIndex(IndexType::Light, DuplicateHandling::PickOne));
 }
diff --git a/clang-tools-extra/clangd/index/FileIndex.h b/clang-tools-extra/clangd/index/FileIndex.h
index d9bee8885d1..87abd0089c0 100644
--- a/clang-tools-extra/clangd/index/FileIndex.h
+++ b/clang-tools-extra/clangd/index/FileIndex.h
@@ -60,21 +60,29 @@ class FileSymbols {
 public:
   /// Updates all symbols and refs in a file.
   /// If either is nullptr, corresponding data for \p Path will be removed.
+  /// If CountReferences is true, \p Refs will be used for counting References
+  /// during merging.
   void update(PathRef Path, std::unique_ptr<SymbolSlab> Slab,
-              std::unique_ptr<RefSlab> Refs);
+              std::unique_ptr<RefSlab> Refs, bool CountReferences);
 
-  // The index keeps the symbols alive.
+  /// The index keeps the symbols alive.
+  /// Will count Symbol::References based on number of references in the main
+  /// files, while building the index with DuplicateHandling::Merge option.
   std::unique_ptr<SymbolIndex>
   buildIndex(IndexType,
              DuplicateHandling DuplicateHandle = DuplicateHandling::PickOne);
 
 private:
+  struct RefSlabAndCountReferences {
+    std::shared_ptr<RefSlab> Slab;
+    bool CountReferences = false;
+  };
   mutable std::mutex Mutex;
 
   /// Stores the latest symbol snapshots for all active files.
   llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSymbols;
   /// Stores the latest ref snapshots for all active files.
-  llvm::StringMap<std::shared_ptr<RefSlab>> FileToRefs;
+  llvm::StringMap<RefSlabAndCountReferences> FileToRefs;
 };
 
 /// This manages symbols from files and an in-memory index on all symbols.
diff --git a/clang-tools-extra/clangd/index/IndexAction.cpp b/clang-tools-extra/clangd/index/IndexAction.cpp
index 120f6b0bae5..c4f553bfb61 100644
--- a/clang-tools-extra/clangd/index/IndexAction.cpp
+++ b/clang-tools-extra/clangd/index/IndexAction.cpp
@@ -186,7 +186,6 @@ std::unique_ptr<FrontendAction> createStaticIndexingAction(
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   Opts.CollectIncludePath = true;
-  Opts.CountReferences = true;
   if (Opts.Origin == SymbolOrigin::Unknown)
     Opts.Origin = SymbolOrigin::Static;
   Opts.StoreAllDocumentation = false;
diff --git a/clang-tools-extra/clangd/indexer/IndexerMain.cpp b/clang-tools-extra/clangd/indexer/IndexerMain.cpp
index d012825101b..530c8891a60 100644
--- a/clang-tools-extra/clangd/indexer/IndexerMain.cpp
+++ b/clang-tools-extra/clangd/indexer/IndexerMain.cpp
@@ -41,6 +41,7 @@ public:
 
   clang::FrontendAction *create() override {
     SymbolCollector::Options Opts;
+    Opts.CountReferences = true;
     return createStaticIndexingAction(
                Opts,
                [&](SymbolSlab S) {
diff --git a/clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp b/clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp
index 86f87009852..df98ee8e507 100644
--- a/clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp
+++ b/clang-tools-extra/clangd/unittests/BackgroundIndexTests.cpp
@@ -32,6 +32,7 @@ MATCHER(EmptyIncludeNode, "") {
   return !arg.IsTU && !arg.URI.empty() && arg.Digest == FileDigest{{0}} &&
          arg.DirectIncludes.empty();
 }
+MATCHER_P(NumReferences, N, "") { return arg.References == N; }
 
 class MemoryShardStorage : public BackgroundIndexStorage {
   mutable std::mutex StorageMu;
@@ -112,6 +113,9 @@ TEST_F(BackgroundIndexTest, IndexTwoFiles) {
       #include "A.h"
       void f_b() {
         (void)common;
+        (void)common;
+        (void)common;
+        (void)common;
       })cpp";
   llvm::StringMap<std::string> Storage;
   size_t CacheHits = 0;
@@ -127,20 +131,25 @@ TEST_F(BackgroundIndexTest, IndexTwoFiles) {
   CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
 
   ASSERT_TRUE(Idx.blockUntilIdleForTest());
-  EXPECT_THAT(
-      runFuzzyFind(Idx, ""),
-      UnorderedElementsAre(Named("common"), Named("A_CC"), Named("g"),
-                           AllOf(Named("f_b"), Declared(), Not(Defined()))));
+  EXPECT_THAT(runFuzzyFind(Idx, ""),
+              UnorderedElementsAre(AllOf(Named("common"), NumReferences(1U)),
+                                   AllOf(Named("A_CC"), NumReferences(0U)),
+                                   AllOf(Named("g"), NumReferences(0U)),
+                                   AllOf(Named("f_b"), Declared(),
+                                         Not(Defined()), NumReferences(0U))));
 
   Cmd.Filename = testPath("root/B.cc");
   Cmd.CommandLine = {"clang++", Cmd.Filename};
-  CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
+  CDB.setCompileCommand(testPath("root/B.cc"), Cmd);
 
   ASSERT_TRUE(Idx.blockUntilIdleForTest());
   // B_CC is dropped as we don't collect symbols from A.h in this compilation.
   EXPECT_THAT(runFuzzyFind(Idx, ""),
-              UnorderedElementsAre(Named("common"), Named("A_CC"), Named("g"),
-                                   AllOf(Named("f_b"), Declared(), Defined())));
+              UnorderedElementsAre(AllOf(Named("common"), NumReferences(5U)),
+                                   AllOf(Named("A_CC"), NumReferences(0U)),
+                                   AllOf(Named("g"), NumReferences(0U)),
+                                   AllOf(Named("f_b"), Declared(), Defined(),
+                                         NumReferences(1U))));
 
   auto Syms = runFuzzyFind(Idx, "common");
   EXPECT_THAT(Syms, UnorderedElementsAre(Named("common")));
@@ -148,6 +157,9 @@ TEST_F(BackgroundIndexTest, IndexTwoFiles) {
   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")}));
 }
 
diff --git a/clang-tools-extra/clangd/unittests/FileIndexTests.cpp b/clang-tools-extra/clangd/unittests/FileIndexTests.cpp
index 142e25540e6..781f5313fbf 100644
--- a/clang-tools-extra/clangd/unittests/FileIndexTests.cpp
+++ b/clang-tools-extra/clangd/unittests/FileIndexTests.cpp
@@ -45,6 +45,7 @@ MATCHER_P(DefURI, U, "") {
   return llvm::StringRef(arg.Definition.FileURI) == U;
 }
 MATCHER_P(QName, N, "") { return (arg.Scope + arg.Name).str() == N; }
+MATCHER_P(NumReferences, N, "") { return arg.References == N; }
 
 namespace clang {
 namespace clangd {
@@ -81,7 +82,7 @@ TEST(FileSymbolsTest, UpdateAndGet) {
   FileSymbols FS;
   EXPECT_THAT(runFuzzyFind(*FS.buildIndex(IndexType::Light), ""), IsEmpty());
 
-  FS.update("f1", numSlab(1, 3), refSlab(SymbolID("1"), "f1.cc"));
+  FS.update("f1", numSlab(1, 3), refSlab(SymbolID("1"), "f1.cc"), false);
   EXPECT_THAT(runFuzzyFind(*FS.buildIndex(IndexType::Light), ""),
               UnorderedElementsAre(QName("1"), QName("2"), QName("3")));
   EXPECT_THAT(getRefs(*FS.buildIndex(IndexType::Light), SymbolID("1")),
@@ -90,8 +91,8 @@ TEST(FileSymbolsTest, UpdateAndGet) {
 
 TEST(FileSymbolsTest, Overlap) {
   FileSymbols FS;
-  FS.update("f1", numSlab(1, 3), nullptr);
-  FS.update("f2", numSlab(3, 5), nullptr);
+  FS.update("f1", numSlab(1, 3), nullptr, false);
+  FS.update("f2", numSlab(3, 5), nullptr, false);
   for (auto Type : {IndexType::Light, IndexType::Heavy})
     EXPECT_THAT(runFuzzyFind(*FS.buildIndex(Type), ""),
                 UnorderedElementsAre(QName("1"), QName("2"), QName("3"),
@@ -110,8 +111,8 @@ TEST(FileSymbolsTest, MergeOverlap) {
   auto X2 = symbol("x");
   X2.Definition.FileURI = "file:///x2";
 
-  FS.update("f1", OneSymboSlab(X1), nullptr);
-  FS.update("f2", OneSymboSlab(X2), nullptr);
+  FS.update("f1", OneSymboSlab(X1), nullptr, false);
+  FS.update("f2", OneSymboSlab(X2), nullptr, false);
   for (auto Type : {IndexType::Light, IndexType::Heavy})
     EXPECT_THAT(
         runFuzzyFind(*FS.buildIndex(Type, DuplicateHandling::Merge), "x"),
@@ -123,14 +124,14 @@ TEST(FileSymbolsTest, SnapshotAliveAfterRemove) {
   FileSymbols FS;
 
   SymbolID ID("1");
-  FS.update("f1", numSlab(1, 3), refSlab(ID, "f1.cc"));
+  FS.update("f1", numSlab(1, 3), refSlab(ID, "f1.cc"), false);
 
   auto Symbols = FS.buildIndex(IndexType::Light);
   EXPECT_THAT(runFuzzyFind(*Symbols, ""),
               UnorderedElementsAre(QName("1"), QName("2"), QName("3")));
   EXPECT_THAT(getRefs(*Symbols, ID), RefsAre({FileURI("f1.cc")}));
 
-  FS.update("f1", nullptr, nullptr);
+  FS.update("f1", nullptr, nullptr, false);
   auto Empty = FS.buildIndex(IndexType::Light);
   EXPECT_THAT(runFuzzyFind(*Empty, ""), IsEmpty());
   EXPECT_THAT(getRefs(*Empty, ID), ElementsAre());
@@ -366,6 +367,33 @@ TEST(FileIndexTest, ReferencesInMainFileWithPreamble) {
               RefsAre({RefRange(Main.range())}));
 }
 
+TEST(FileSymbolsTest, CountReferencesNoRefSlabs) {
+  FileSymbols FS;
+  FS.update("f1", numSlab(1, 3), nullptr, true);
+  FS.update("f2", numSlab(1, 3), nullptr, false);
+  EXPECT_THAT(
+      runFuzzyFind(*FS.buildIndex(IndexType::Light, DuplicateHandling::Merge),
+                   ""),
+      UnorderedElementsAre(AllOf(QName("1"), NumReferences(0u)),
+                           AllOf(QName("2"), NumReferences(0u)),
+                           AllOf(QName("3"), NumReferences(0u))));
+}
+
+TEST(FileSymbolsTest, CountReferencesWithRefSlabs) {
+  FileSymbols FS;
+  FS.update("f1cpp", numSlab(1, 3), refSlab(SymbolID("1"), "f1.cpp"), true);
+  FS.update("f1h", numSlab(1, 3), refSlab(SymbolID("1"), "f1.h"), false);
+  FS.update("f2cpp", numSlab(1, 3), refSlab(SymbolID("2"), "f2.cpp"), true);
+  FS.update("f2h", numSlab(1, 3), refSlab(SymbolID("2"), "f2.h"), false);
+  FS.update("f3cpp", numSlab(1, 3), refSlab(SymbolID("3"), "f3.cpp"), true);
+  FS.update("f3h", numSlab(1, 3), refSlab(SymbolID("3"), "f3.h"), false);
+  EXPECT_THAT(
+      runFuzzyFind(*FS.buildIndex(IndexType::Light, DuplicateHandling::Merge),
+                   ""),
+      UnorderedElementsAre(AllOf(QName("1"), NumReferences(1u)),
+                           AllOf(QName("2"), NumReferences(1u)),
+                           AllOf(QName("3"), NumReferences(1u))));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang