1 files changed, 35 insertions, 7 deletions

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