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

[kadircet]

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.

[ioeric]

I'm not sure if FileIndex is the correct layer to make decision about how References is calculated. Currently, we have two use cases in clangd 1) one slab per TU and 2) one slab per file. It seems to me that we would want different strategies for the two use cases, so it might make sense to make this configurable in FileIndex. And in one case, we have References ~= # of TUs, and in the other case, we would have References ~= # of files. This can over-count References in the second case. It might be fine as an approximation (if there is no better solution), but we should definitely document it (e.g. in BackgroundIndex).

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

50

this is a bit of a code smell. FileIndex is not supposed to know anything about clangd-indexer etc.

136

note that references might not always exist. SymbolCollector can be set up to not collect references.

[kadircet]

In D59481#1432881, @ioeric wrote:

I'm not sure if FileIndex is the correct layer to make decision about how References is calculated. Currently, we have two use cases in clangd 1) one slab per TU and 2) one slab per file. It seems to me that we would want different strategies for the two use cases, so it might make sense to make this configurable in FileIndex. And in one case, we have References ~= # of TUs, and in the other case, we would have References ~= # of files.

References ~= # of TUs case has nothing to do with FileIndex actually, it only happens if references are counted beforehand(which is currently only happening in clangd-indexer).
FileIndex has two ways to handle duplicates during an index merge:

• pickone which doesn't count references at all (this is the oneslab per TU case)
• merge which counts references as number of files this symbol has occured (this is the oneslab per File case)

This can over-count References in the second case. It might be fine as an approximation (if there is no better solution), but we should definitely document it (e.g. in BackgroundIndex).

I don't follow why this can over-count, FileIndex keeps only one RefSlab per each file, so we can't over-count? Did you mean it will be more than #TUs ?

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

50

I've actually just left the comment for review so that we can decide what to do about discrepancy(with my reasoning). It is not like FileIndex is somehow tied to clangd-indexer now?

136

I thought we wouldn't make any promises about References in such a case, do we?

[ioeric]

I don't follow why this can over-count, FileIndex keeps only one RefSlab per each file, so we can't over-count? Did you mean it will be more than #TUs ?

Yes. This is how Symbol::References is described in its documentation. If we want to change/expand the semantic, we need to make it clear in the documentation as well.

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

136

Is this documented somewhere?

References and RefSlab have been two independent things. References existed before RefSlab. Now might be a good time to unify them, but I think we should think about how this is reflected in their definitions (documentation etc) and other producers and consumers of the structures.

[kadircet]

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

136

Both of them are only produced by SymbolCollector, which has an option CountReferences with a comment saying // Populate the Symbol.References field..
Unfortunately that's not what it does, instead it always sets Symol.References field to 1 for every symbol collected during indexing of a single TU.
Then it is up to the consumers to merge those 1s. So even though we say:

/// The number of translation units that reference this symbol from their main
/// file. This number is only meaningful if aggregated in an index.

in the comments of Symbol::References it is totally up-to-the consumer who performs the merging.

The option controls whether Symbol Occurences should be collected or not.

I see two possible options:

• Un-tie SymbolCollector from Symbol::References and rather define it as:

/// This field is for the convenience of an aggregated symbol index

• Rather change Index structure's(consumers of this information) to store(and build) this information internally if they plan to make use of it ?

WDYT?

[ioeric]

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

33

nit: s/every file/every slab/?

There is no "file" in this context.

46

I think we can trigger this assertion with an incomplete background index. For example, we've seen references of a symbol in some files but we've not parsed the file that contains the decls.

50

the comment is useful. I just think it's in the wrong place. If we define the reference counting behavior for FileSymbols, FileSymbolswouldn't need to know about clangd-indexer or background-index. This comment can then live in background index instead.

136

Unfortunately that's not what it does, instead it always sets Symol.References field to 1 for every symbol collected during indexing of a single TU.

Looking at the code, I think Symol.References is indeed only populated (i.e. set to 1) when CountReferences is enabled. Honestly, we should probably get rid of CountReferences; I don't see a reason not to count references.

But the statement count one per TU still stands in SymbolCollector, as it always only handles one TU at a time.

Actually, all I'm asking in the initial comment is to define the References counting behavior when there is no reference ;) I think a reasonable solution is to keep the references untouched when there is no reference, so that References aggregation via mergeSymbol can still work (for one-TU-per-Slab use cases).

And we should probably also document the new behavior e.g. DuplicateHandling in FileIndex.h seems to be a good place.

148–149

this FIXME can be removed?

159

any reason not to count references for PickOne?

[ilya-biryukov]

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

136

Honestly, we should probably get rid of CountReferences; I don't see a reason not to count references

We introduced it to avoid getting partial counts in dynamic index. In the absence of the static index, they would skew ranking towards symbols used inside the files you have open, which is not really a good signal in the first place.

I'm not sure if that actually hurts much in practice, though. It could very possibly be the case that the added complexity isn't worth it.

[kadircet]

• Address comments

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

46

You are right, thanks for pointing this out!

50

Moved into index/background.h

159

Because PickOne is not populating SymsStorage and is merely passing along pointers to SymbolSlabs. Counting references in that case would imply also making a new copy per each unique symbol.

Do you think it is worth it?

[ilya-biryukov]

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

113 ↗

(On Diff #191844)

We should agree on and stick to a single behavior in both cases.
I suggest keeping the current behavior for now (translation units). Why can't we implement that?

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

52

The comment suggests there's something cheesy going on.
Why would FileSymbols recompute the number of references? Can we avoid this complicated logic?

[kadircet]

It has been a long time since I've proposed that change and even I forgot some of the high level details. Therefore, I wanted to sum up the state again so that we can decide on how to move forward.

Currently we have two different on-disk formats for index in clangd:

• Static index:
  • Merged single file, since merging is done before persistence, the data on-disk contains correct reference counts.
  • Produced by running SymbolCollector on each TU. Collector doesn't keep state between different TUs, merging is done on the caller-site.
  • No need to count references when loading, doesn't interact with FileSymbols.

• Background index:
  • Sharded multiple files.
  • Produced by running SymbolCollector on each TU, but instead of merging we separate those symbols into distinct files. Merging is done later on within FileSymbols which has no information about TUs.

FileSymbols is currently being used by BackgroundIndex and FileIndex(Dynamic Idx):

• In the case of BackgroundIndex, a symbol occurs at most twice(once in the header declaring it and once in the implementation file, if they are distinct), therefore when merging is done symbol references doesn't go above that.
• FileIndex doesn't perform de-duplication before-hand therefore it can perform reference counting while performing the merge. But currently it doesn't perform any merging.

As a result, changing FileSymbols would only effect BackgroundIndex and nothing else. Since merging occurs using the information in RefSlabs this can also be extended to benefit FileIndex and will most definitely be used
by any sort of indexing that performs merging over multiple files.

For unification purposes I would rather delete the References counting logic in SymbolCollector and perform it only at mergers(clangd-indexer and FileSymbols)

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

113 ↗

(On Diff #191844)

Because FileSymbols only knows about file names and has no idea about whether a given file is TU or not. We might add some heuristics on the file extensions, or change SymbolCollector to count number of files(this would require maintaining a cache to de-duplicate references coming from the same files but originating from a different TU).

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

52

This was the idea behind https://github.com/clangd/clangd/issues/23. Please see the main comment for a higher level discussion.

[ilya-biryukov]

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

113 ↗

(On Diff #191844)

Update from the offline conversation:
It should be possible to produce the same counts as static-indexer, will wait for an updated revision.

[kadircet]

• Change logic to count references from only main files

[ilya-biryukov]

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

482

NIT: maybe initialize with =false to avoid potential UB.
Digest seems uninitialized too, could also ={} for it.

Sorry if this was discussed before and I'm repeating myself, I vaguely remember something similar but not sure what the conclusion was.

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

120

Let's use a named struct here? .second.second is super-hard to read

129

Could we avoid changing the merge logic?
I.e. the proposal is to keep the merging code the same and add a loop that counts references at the end.

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

60

NIT: this comment seems more appropriate at buildIndex, maybe move it there?
(It has the DuplicateHandle parameter, so it's closer to the context.

63

Maybe call the parameter CountReferences?
It would narrow down what the code can do with it and make the code easier to understand

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

49

NIT: to avoid warnings, we could use the corresponding suffix at the callsite (e.g. 1u) instead of casts
It's probably ok either way here, but having casts can lead to surprises if one passes negative numbers there

[kadircet]

• Address comments

[ilya-biryukov]

LGTM

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

101–106

NIT: maybe simplify to the following?

if (!Refs) {
  FileToRefs.erase(Path);
  return;
}
...

up to you.

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

63

NIT: use \p Refs

[kadircet]

• Address comments

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

if (IndexStorage) {
       // Skip if file is already up to date.
       auto DigestIt = IndexedFileDigests.try_emplace(Path);
       if (!DigestIt.second && DigestIt.first->second == Hash)
         continue;
       DigestIt.first->second = Hash;
       // 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);
     }
   }
 }
 
 void BackgroundIndex::buildIndex() {
   assert(BuildIndexPeriodMs > 0);
   while (true) {
     {

 
 std::vector<BackgroundIndex::Source>
 BackgroundIndex::loadShard(const tooling::CompileCommand &Cmd,
                            BackgroundIndexStorage *IndexStorage,
                            llvm::StringSet<> &LoadedShards) {
   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
   // paths.
   llvm::StringSet<> InQueue;
   auto FS = FSProvider.getFileSystem();
   // Dependencies of this TU, paired with the information about whether they
   // need to be re-indexed or not.
for (const auto &I : *Shard->Sources) {
         continue;
 
       // We found source file info for current dependency.
       assert(I.getValue().Digest != FileDigest{{0}} && "Digest is empty?");
       ShardInfo SI;
       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.
       auto Buf = FS->getBufferForFile(CurDependency.Path);
       if (!Buf) {
         elog("Couldn't get buffer for file: {0}: {1}", CurDependency.Path,
              Buf.getError().message());
         continue;
for (const ShardInfo &SI : IntermediateSymbols) {
       auto SS =
           SI.Shard->Symbols
               ? llvm::make_unique<SymbolSlab>(std::move(*SI.Shard->Symbols))
               : nullptr;
       auto RS = SI.Shard->Refs
                     ? 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);
     }
   }
 
   return Dependencies;
 }
 
 // Goes over each changed file and loads them from index. Returns the list of
 // TUs that had out-of-date/no shards.

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

 /// managed with shared pointers that are shared between this class and the
 /// users. For each file, this class only stores a pointer pointing to the
 /// newest snapshot, and an outdated snapshot is deleted by the last owner of
 /// the snapshot, either this class or the symbol index.
 ///
 /// The snapshot semantics keeps critical sections minimal since we only need
 /// locking when we swap or obtain references to snapshots.
 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.
 /// FIXME: Expose an interface to remove files that are closed.
 class FileIndex : public MergedIndex {
 public:
   FileIndex(bool UseDex = true);
 

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

 #include "llvm/ADT/StringRef.h"
 #include <memory>
 
 namespace clang {
 namespace clangd {
 
 static std::pair<SymbolSlab, RefSlab>
 indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
              llvm::ArrayRef<Decl *> DeclsToIndex,
              const CanonicalIncludes &Includes, bool IsIndexMainAST) {
   SymbolCollector::Options CollectorOpts;
   CollectorOpts.CollectIncludePath = true;
   CollectorOpts.Includes = &Includes;
   CollectorOpts.CountReferences = false;
   CollectorOpts.Origin = SymbolOrigin::Dynamic;
 
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
   if (IsIndexMainAST) {
     // We only collect refs when indexing main AST.
     CollectorOpts.RefFilter = RefKind::All;
     // Comments for main file can always be obtained from sema, do not store
     // them in the index.
     CollectorOpts.StoreAllDocumentation = false;
   } else {
     IndexOpts.IndexMacrosInPreprocessor = true;
     CollectorOpts.CollectMacro = true;
     CollectorOpts.StoreAllDocumentation = true;
   }
 
   SymbolCollector Collector(std::move(CollectorOpts));
   Collector.setPreprocessor(PP);
   index::indexTopLevelDecls(AST, *PP, DeclsToIndex, Collector, IndexOpts);
std::vector<Decl *> DeclsToIndex(
       AST.getTranslationUnitDecl()->decls().begin(),
       AST.getTranslationUnitDecl()->decls().end());
   return indexSymbols(AST, std::move(PP), DeclsToIndex, Includes,
                       /*IsIndexMainAST=*/false)
       .first;
 }
 
 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;
   switch (DuplicateHandle) {
   case DuplicateHandling::Merge: {
     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;
   }
   }
 
   std::vector<Ref> RefsStorage; // Contiguous ranges for each SymbolID.
   llvm::DenseMap<SymbolID, llvm::ArrayRef<Ref>> AllRefs;
   {
     llvm::DenseMap<SymbolID, llvm::SmallVector<Ref, 4>> MergedRefs;
FileIndex::FileIndex(bool UseDex)
     : MergedIndex(&MainFileIndex, &PreambleIndex), UseDex(UseDex),
       PreambleIndex(llvm::make_unique<MemIndex>()),
       MainFileIndex(llvm::make_unique<MemIndex>()) {}
 
 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));
 }
 
 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));
 }
 
 } // namespace clangd
 } // namespace clang

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

std::unique_ptr<FrontendAction> createStaticIndexingAction(
     SymbolCollector::Options Opts,
     std::function<void(SymbolSlab)> SymbolsCallback,
     std::function<void(RefSlab)> RefsCallback,
     std::function<void(IncludeGraph)> IncludeGraphCallback) {
   index::IndexingOptions IndexOpts;
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   Opts.CollectIncludePath = true;
-  Opts.CountReferences = true;
   if (Opts.Origin == SymbolOrigin::Unknown)
     Opts.Origin = SymbolOrigin::Static;
   Opts.StoreAllDocumentation = false;
   if (RefsCallback != nullptr) {
     Opts.RefFilter = RefKind::All;
     Opts.RefsInHeaders = true;
   }
   auto Includes = llvm::make_unique<CanonicalIncludes>();

clang-tools-extra/clangd/indexer/IndexerMain.cpp

Format("format", llvm::cl::desc("Format of the index to be written"),
            llvm::cl::init(IndexFileFormat::RIFF));
 
 class IndexActionFactory : public tooling::FrontendActionFactory {
 public:
   IndexActionFactory(IndexFileIn &Result) : Result(Result) {}
 
   clang::FrontendAction *create() override {
     SymbolCollector::Options Opts;
+    Opts.CountReferences = true;
     return createStaticIndexingAction(
                Opts,
                [&](SymbolSlab S) {
                  // Merge as we go.
                  std::lock_guard<std::mutex> Lock(SymbolsMu);
                  for (const auto &Sym : S) {
                    if (const auto *Existing = Symbols.find(Sym.ID))
                      Symbols.insert(mergeSymbol(*Existing, Sym));

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

 RefsAre(std::vector<::testing::Matcher<Ref>> Matchers) {
   return ElementsAre(::testing::Pair(_, UnorderedElementsAreArray(Matchers)));
 }
 // URI cannot be empty since it references keys in the IncludeGraph.
 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;
   llvm::StringMap<std::string> &Storage;
   size_t &CacheHits;
 
 public:
   MemoryShardStorage(llvm::StringMap<std::string> &Storage, size_t &CacheHits)
TEST_F(BackgroundIndexTest, IndexTwoFiles) {
   FS.Files[testPath("root/A.cc")] =
       "#include \"A.h\"\nvoid g() { (void)common; }";
   FS.Files[testPath("root/B.cc")] =
       R"cpp(
       #define A 0
       #include "A.h"
       void f_b() {
         (void)common;
+        (void)common;
+        (void)common;
+        (void)common;
       })cpp";
   llvm::StringMap<std::string> Storage;
   size_t CacheHits = 0;
   MemoryShardStorage MSS(Storage, CacheHits);
   OverlayCDB CDB(/*Base=*/nullptr);
   BackgroundIndex Idx(Context::empty(), FS, CDB,
                       [&](llvm::StringRef) { return &MSS; });
 
   tooling::CompileCommand Cmd;
   Cmd.Filename = testPath("root/A.cc");
   Cmd.Directory = testPath("root");
   Cmd.CommandLine = {"clang++", "-DA=1", testPath("root/A.cc")};
   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")));
   auto Common = *Syms.begin();
   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")}));
 }
 
 TEST_F(BackgroundIndexTest, ShardStorageTest) {
   MockFSProvider FS;
   FS.Files[testPath("root/A.h")] = R"cpp(
       void common();
       void f_b();

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

 MATCHER_P(FileURI, F, "") { return llvm::StringRef(arg.Location.FileURI) == F; }
 MATCHER_P(DeclURI, U, "") {
   return llvm::StringRef(arg.CanonicalDeclaration.FileURI) == U;
 }
 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 {
 namespace {
 ::testing::Matcher<const RefSlab &>
 RefsAre(std::vector<::testing::Matcher<Ref>> Matchers) {
   return ElementsAre(::testing::Pair(_, UnorderedElementsAreArray(Matchers)));
 }
 
std::unique_ptr<RefSlab> refSlab(const SymbolID &ID, const char *Path) {
   Slab.insert(ID, R);
   return llvm::make_unique<RefSlab>(std::move(Slab).build());
 }
 
 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")),
               RefsAre({FileURI("f1.cc")}));
 }
 
 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"),
                                      QName("4"), QName("5")));
 }
 
 TEST(FileSymbolsTest, MergeOverlap) {
   FileSymbols FS;
   auto OneSymboSlab = [](Symbol Sym) {
     SymbolSlab::Builder S;
     S.insert(Sym);
     return llvm::make_unique<SymbolSlab>(std::move(S).build());
   };
   auto X1 = symbol("x");
   X1.CanonicalDeclaration.FileURI = "file:///x1";
   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"),
         UnorderedElementsAre(
             AllOf(QName("x"), DeclURI("file:///x1"), DefURI("file:///x2"))));
 }
 
 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());
 
   EXPECT_THAT(runFuzzyFind(*Symbols, ""),
               UnorderedElementsAre(QName("1"), QName("2"), QName("3")));
   EXPECT_THAT(getRefs(*Symbols, ID), RefsAre({FileURI("f1.cc")}));
 }
TEST(FileIndexTest, ReferencesInMainFileWithPreamble) {
   Index.updateMain(testPath(TU.Filename), AST);
 
   // Expect to see references in main file, references in headers are excluded
   // because we only index main AST.
   EXPECT_THAT(getRefs(Index, findSymbol(TU.headerSymbols(), "Foo").ID),
               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