1 files changed, 81 insertions, 1 deletions

diff --git a/clang-tools-extra/unittests/clangd/BackgroundIndexTests.cpp b/clang-tools-extra/unittests/clangd/BackgroundIndexTests.cpp
index c6f401d2740..6746e800ea1 100644
--- a/clang-tools-extra/unittests/clangd/BackgroundIndexTests.cpp
+++ b/clang-tools-extra/unittests/clangd/BackgroundIndexTests.cpp
@@ -1,6 +1,7 @@
 #include "SyncAPI.h"
 #include "TestFS.h"
 #include "index/Background.h"
+#include "llvm/Support/ScopedPrinter.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
@@ -24,6 +25,37 @@ RefsAre(std::vector<testing::Matcher<Ref>> Matchers) {
   return ElementsAre(testing::Pair(_, UnorderedElementsAreArray(Matchers)));
 }
 
+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)
+      : Storage(Storage), CacheHits(CacheHits) {}
+  llvm::Error storeShard(llvm::StringRef ShardIdentifier,
+                         IndexFileOut Shard) const override {
+    std::lock_guard<std::mutex> Lock(StorageMu);
+    Storage[ShardIdentifier] = llvm::to_string(Shard);
+    return llvm::Error::success();
+  }
+  std::unique_ptr<IndexFileIn>
+  loadShard(llvm::StringRef ShardIdentifier) const override {
+    std::lock_guard<std::mutex> Lock(StorageMu);
+    if (Storage.find(ShardIdentifier) == Storage.end()) {
+      return nullptr;
+    }
+    auto IndexFile = readIndexFile(Storage[ShardIdentifier]);
+    if (!IndexFile) {
+      ADD_FAILURE() << "Error while reading " << ShardIdentifier << ':'
+                    << IndexFile.takeError();
+      return nullptr;
+    }
+    CacheHits++;
+    return llvm::make_unique<IndexFileIn>(std::move(*IndexFile));
+  }
+};
+
 TEST(BackgroundIndexTest, IndexTwoFiles) {
   MockFSProvider FS;
   // a.h yields different symbols when included by A.cc vs B.cc.
@@ -45,7 +77,11 @@ TEST(BackgroundIndexTest, IndexTwoFiles) {
       void f_b() {
         (void)common;
       })cpp";
-  BackgroundIndex Idx(Context::empty(), "", FS, /*URISchemes=*/{"unittest"});
+  llvm::StringMap<std::string> Storage;
+  size_t CacheHits = 0;
+  MemoryShardStorage MSS(Storage, CacheHits);
+  BackgroundIndex Idx(Context::empty(), "", FS, /*URISchemes=*/{"unittest"},
+                      [&](llvm::StringRef) { return &MSS; });
 
   tooling::CompileCommand Cmd;
   Cmd.Filename = testPath("root/A.cc");
@@ -78,5 +114,49 @@ TEST(BackgroundIndexTest, IndexTwoFiles) {
                        FileURI("unittest:///root/B.cc")}));
 }
 
+TEST(BackgroundIndexTest, ShardStorageWriteTest) {
+  MockFSProvider FS;
+  FS.Files[testPath("root/A.h")] = R"cpp(
+      void common();
+      void f_b();
+      class A_CC {};
+      )cpp";
+  FS.Files[testPath("root/A.cc")] =
+      "#include \"A.h\"\nvoid g() { (void)common; }";
+
+  llvm::StringMap<std::string> Storage;
+  size_t CacheHits = 0;
+  MemoryShardStorage MSS(Storage, CacheHits);
+
+  tooling::CompileCommand Cmd;
+  Cmd.Filename = testPath("root/A.cc");
+  Cmd.Directory = testPath("root");
+  Cmd.CommandLine = {"clang++", testPath("root/A.cc")};
+  // Check nothing is loaded from Storage, but A.cc and A.h has been stored.
+  {
+    BackgroundIndex Idx(Context::empty(), "", FS, /*URISchemes=*/{"unittest"},
+                        [&](llvm::StringRef) { return &MSS; });
+    Idx.enqueue(testPath("root"), Cmd);
+    Idx.blockUntilIdleForTest();
+  }
+  EXPECT_EQ(CacheHits, 0U);
+  EXPECT_EQ(Storage.size(), 2U);
+
+  auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
+  EXPECT_NE(ShardHeader, nullptr);
+  EXPECT_THAT(
+      *ShardHeader->Symbols,
+      UnorderedElementsAre(Named("common"), Named("A_CC"),
+                           AllOf(Named("f_b"), Declared(), Not(Defined()))));
+  for (const auto &Ref : *ShardHeader->Refs)
+    EXPECT_THAT(Ref.second,
+                UnorderedElementsAre(FileURI("unittest:///root/A.h")));
+
+  auto ShardSource = MSS.loadShard(testPath("root/A.cc"));
+  EXPECT_NE(ShardSource, nullptr);
+  EXPECT_THAT(*ShardSource->Symbols, UnorderedElementsAre());
+  EXPECT_THAT(*ShardSource->Refs, RefsAre({FileURI("unittest:///root/A.cc")}));
+}
+
 } // namespace clangd
 } // namespace clang