[clangd] Refactor background-index shard loading

[sammccall]

So I've stared at this refactoring for a while, and I still don't totally get it.
It seems like a class that really wants to be a function. It's a complicated function though - maybe a separate cpp file does make sense.

We discussed moving ShardVersions and Writes into the loader, splitting ShardVersions from IndexSymbols seems pretty odd though.

What's the advantage of making this stateful?

[kadircet]

• As discussed offline, moved the "staleness" detection back to background index
• A renaming within Rebuilder, "LoadedTUs -> LoadedShards"

[sammccall]

Refactoring generally looks good.
You're replacing a lot of documented code with new undocumented code, can we add some high-level comments?

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

127

nit: unclear what's changed with respect to what.

I'd suggest shardIsStale() and swap the param order?

130

While here: this message lacks context and is unnecessarily jargony: "Background-indexing: couldn't read {0} to validate stored index"?

424

move this outside the loop and use a counter? it's an extra lock/unlock for no reason

425

This function now has almost no comments, can you explain the major steps?

443

nit: again, this patch is adding Path/PathRef aliases to parts of the code that don't use them

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

169

and their dependencies

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

78

you can't call this an (opaque?) ShardIdentifier and then use it as a path - maybe SourceFile?

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

2

header is out of date. The new file/structs/functions need documentation

26

LoadedShard?

31

IndexFileIn is just a struct - optional<>? (And need a comment explaining when it can be null)

37

no, they're not :-)

we talked about making the values refs to ShardInfo.AbsolutePath.

38

why are these outside the ShardInfo?

43

clangd::load seems like an insufficiently precise name for this function

loadIndexShards?

[kadircet]

• Address comments

[kadircet]

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

443

I believe they represent the intention better then string/StringRefthat was the reason I was using them.

[sammccall]

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

130

{0 is missing its }

425

// We'll accept data from stale shards, but ensure the files get reindexed soon.

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

94

This handles only one main file, I can't see where you skip loading shards for headers that were loaded for a previous file

105

ShardIdentifier still here

127

I don't think this works, the storage for these StringRefs is the LoadedShards you just moved from.
(In practice, maybe you get away with it if the strings are longer than the SSO?)

[kadircet]

• Address comments

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

94

The real loading happens down the line at const CachedShard &CS = loadShard(ShardIdentifier, Storage);, which returns shard from cache if it was already loaded.

127

Yes you are right, I don't see a good way out of it with current data model. I suppose we can keep an "intern table :)" for the strings, but then LoadedShards would be meaningless without that intern table so we would end up returning another struct from loadIndexShards that will wrap std::vector<LoadedShards> and the intern table(but we tried to mitigate from that struct in previous iteration).

So instead I am changing the model to return only the first translation unit and actually this makes a lot of sense when we think about shards for standard library, in a code base like chromium vector could easily have ~40k dependent TUs and even if we just returned a vector<StringRef> it would still be around 640kB, and we have a lot of standard library headers. this can quite easily add hundreds of MB to clangd's memory footprint while loading shards.

[kadircet]

• Add comments

[sammccall]

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

48

I think we only need to traverse this when *missing* the cache, so it can be a return value of loadShard rather than stored in CachedShard (so CachedShard becomes just LoadedShard)

[sammccall]

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

94

OK, then we still traverse the whole subtree but hit every node.

I think it'd be nicer to avoid the traversal (as noted above) if that works well, up to you.

[kadircet]

• Get rid of CachedShard

clang-tools-extra/clangd/CMakeLists.txt

add_clang_library(clangDaemon
   QueryDriverDatabase.cpp
   Threading.cpp
   Trace.cpp
   TUScheduler.cpp
   URI.cpp
   XRefs.cpp
 
   index/Background.cpp
+  index/BackgroundIndexLoader.cpp
   index/BackgroundIndexStorage.cpp
   index/BackgroundQueue.cpp
   index/BackgroundRebuild.cpp
   index/CanonicalIncludes.cpp
   index/FileIndex.cpp
   index/Index.cpp
   index/IndexAction.cpp
   index/MemIndex.cpp

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

 //===--- Background.h - Build an index in a background thread ----*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_H
 
 #include "Context.h"
 #include "FSProvider.h"
 #include "GlobalCompilationDatabase.h"
+#include "Path.h"
 #include "SourceCode.h"
 #include "Threading.h"
 #include "index/BackgroundRebuild.h"
 #include "index/FileIndex.h"
 #include "index/Index.h"
 #include "index/Serialization.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/StringMap.h"
llvm::Error index(tooling::CompileCommand,
                     BackgroundIndexStorage *IndexStorage);
 
   FileSymbols IndexedSymbols;
   BackgroundIndexRebuilder Rebuilder;
   llvm::StringMap<ShardVersion> ShardVersions; // Key is absolute file path.
   std::mutex ShardVersionsMu;
 
   BackgroundIndexStorage::Factory IndexStorageFactory;
-  struct Source {
-    std::string Path;
-    bool NeedsReIndexing;
-    Source(llvm::StringRef Path, bool NeedsReIndexing)
-        : Path(Path), NeedsReIndexing(NeedsReIndexing) {}
-  };
-  // Loads the shards for a single TU and all of its dependencies. Returns the
-  // list of sources and whether they need to be re-indexed.
-  std::vector<Source> loadShard(const tooling::CompileCommand &Cmd,
-                                BackgroundIndexStorage *IndexStorage,
-                                llvm::StringSet<> &LoadedShards);
-  // Tries to load shards for the ChangedFiles.
+  // Tries to load shards for the MainFiles.
   std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
-  loadShards(std::vector<std::string> ChangedFiles);
+  loadProject(std::vector<std::string> MainFiles);
 
   BackgroundQueue::Task
   changedFilesTask(const std::vector<std::string> &ChangedFiles);
   BackgroundQueue::Task indexFileTask(tooling::CompileCommand Cmd,
                                       BackgroundIndexStorage *Storage);
 
   // from lowest to highest priority
   enum QueuePriority {

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

 //===-- Background.cpp - Build an index in a background thread ------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "index/Background.h"
 #include "ClangdUnit.h"
 #include "Compiler.h"
 #include "Context.h"
+#include "FSProvider.h"
 #include "Headers.h"
 #include "Logger.h"
 #include "Path.h"
 #include "SourceCode.h"
 #include "Symbol.h"
 #include "Threading.h"
 #include "Trace.h"
 #include "URI.h"
+#include "index/BackgroundIndexLoader.h"
 #include "index/FileIndex.h"
 #include "index/IndexAction.h"
 #include "index/MemIndex.h"
 #include "index/Ref.h"
 #include "index/Relation.h"
 #include "index/Serialization.h"
 #include "index/SymbolCollector.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Threading.h"
 
 #include <algorithm>
 #include <atomic>
 #include <chrono>
 #include <condition_variable>
+#include <cstddef>
 #include <memory>
 #include <mutex>
 #include <numeric>
 #include <queue>
 #include <random>
 #include <string>
 #include <thread>
+#include <utility>
+#include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 // Resolves URI to file paths with cache.
 class URIToFileCache {
 public:
if (llvm::sys::path::is_absolute(Cmd.Filename)) {
     AbsolutePath = Cmd.Filename;
   } else {
     AbsolutePath = Cmd.Directory;
     llvm::sys::path::append(AbsolutePath, Cmd.Filename);
     llvm::sys::path::remove_dots(AbsolutePath, true);
   }
   return AbsolutePath;
 }
+
 } // namespace
 
 BackgroundIndex::BackgroundIndex(
     Context BackgroundContext, const FileSystemProvider &FSProvider,
     const GlobalCompilationDatabase &CDB,
     BackgroundIndexStorage::Factory IndexStorageFactory, size_t ThreadPoolSize)
     : SwapIndex(llvm::make_unique<MemIndex>()), FSProvider(FSProvider),
       CDB(CDB), BackgroundContext(std::move(BackgroundContext)),
       Rebuilder(this, &IndexedSymbols),
       IndexStorageFactory(std::move(IndexStorageFactory)),
       CommandsChanged(
           CDB.watch([&](const std::vector<std::string> &ChangedFiles) {

 BackgroundQueue::Task BackgroundIndex::changedFilesTask(
     const std::vector<std::string> &ChangedFiles) {
   BackgroundQueue::Task T([this, ChangedFiles] {
     trace::Span Tracer("BackgroundIndexEnqueue");
     // We're doing this asynchronously, because we'll read shards here too.
     log("Enqueueing {0} commands for indexing", ChangedFiles.size());
     SPAN_ATTACH(Tracer, "files", int64_t(ChangedFiles.size()));
 
-    auto NeedsReIndexing = loadShards(std::move(ChangedFiles));
+    auto NeedsReIndexing = loadProject(std::move(ChangedFiles));
     // Run indexing for files that need to be updated.
     std::shuffle(NeedsReIndexing.begin(), NeedsReIndexing.end(),
                  std::mt19937(std::random_device{}()));
     std::vector<BackgroundQueue::Task> Tasks;
     Tasks.reserve(NeedsReIndexing.size());
     for (auto &Elem : NeedsReIndexing)
       Tasks.push_back(indexFileTask(std::move(Elem.first), Elem.second));
     Queue.append(std::move(Tasks));
llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd,
   }
   update(AbsolutePath, std::move(Index), ShardVersionsSnapshot, IndexStorage,
          HadErrors);
 
   Rebuilder.indexedTU();
   return llvm::Error::success();
 }
 
-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 = {};
-    bool CountReferences = false;
-    bool HadErrors = false;
-  };
-  std::vector<ShardInfo> IntermediateSymbols;
-  // Make sure we don't have duplicate elements in the queue. Keys are absolute
-  // paths.
-  llvm::StringSet<> InQueue;
+// Goes over each changed file and loads them from index. Returns the list of
+// TUs that had out-of-date/no shards.
+std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
+BackgroundIndex::loadProject(std::vector<std::string> MainFiles) {
+  std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
+      NeedsReIndexing;
+
   auto FS = FSProvider.getFileSystem();
-  // Dependencies of this TU, paired with the information about whether they
-  // need to be re-indexed or not.
-  std::vector<Source> Dependencies;
-  std::queue<Source> ToVisit;
-  std::string AbsolutePath = getAbsolutePath(Cmd).str();
-  // Up until we load the shard related to a dependency it needs to be
-  // re-indexed.
-  ToVisit.emplace(AbsolutePath, true);
-  InQueue.insert(AbsolutePath);
-  // Goes over each dependency.
-  while (!ToVisit.empty()) {
-    Dependencies.push_back(std::move(ToVisit.front()));
-    // Dependencies is not modified during the rest of the loop, so it is safe
-    // to keep the reference.
-    auto &CurDependency = Dependencies.back();
-    ToVisit.pop();
-    // If we have already seen this shard before(either loaded or failed) don't
-    // re-try again. Since the information in the shard won't change from one TU
-    // to another.
-    if (!LoadedShards.try_emplace(CurDependency.Path).second) {
-      // If the dependency needs to be re-indexed, first occurence would already
-      // have detected that, so we don't need to issue it again.
-      CurDependency.NeedsReIndexing = false;
-      continue;
-    }
+  BackgroundIndexLoader Loader(FS.get());
 
-    auto Shard = IndexStorage->loadShard(CurDependency.Path);
-    if (!Shard || !Shard->Sources) {
-      // File will be returned as requiring re-indexing to caller.
-      vlog("Failed to load shard: {0}", CurDependency.Path);
-      continue;
-    }
-    // These are the edges in the include graph for current dependency.
-    for (const auto &I : *Shard->Sources) {
-      auto U = URI::parse(I.getKey());
-      if (!U)
-        continue;
-      auto AbsolutePath = URI::resolve(*U, CurDependency.Path);
-      if (!AbsolutePath)
-        continue;
-      // Add file as dependency if haven't seen before.
-      if (InQueue.try_emplace(*AbsolutePath).second)
-        ToVisit.emplace(*AbsolutePath, true);
-      // The node contains symbol information only for current file, the rest is
-      // just edges.
-      if (*AbsolutePath != CurDependency.Path)
-        continue;
+  Rebuilder.startLoading();
+  for (const auto &File : MainFiles) {
+    auto Cmd = CDB.getCompileCommand(File);
+    if (!Cmd)
 
-      // 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().Flags & IncludeGraphNode::SourceFlag::IsTU;
-      SI.HadErrors =
-          I.getValue().Flags & IncludeGraphNode::SourceFlag::HadErrors;
-      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;
-      }
-      // If digests match then dependency doesn't need re-indexing.
-      // FIXME: Also check for dependencies(sources) of this shard and compile
-      // commands for cache invalidation.
-      CurDependency.NeedsReIndexing =
-          digest(Buf->get()->getBuffer()) != I.getValue().Digest;
-    }
+    std::string ProjectRoot;
+    if (auto PI = CDB.getProjectInfo(File))
+      ProjectRoot = std::move(PI->SourceRoot);
+
+    BackgroundIndexStorage *IndexStorage = IndexStorageFactory(ProjectRoot);
+    Path AbsPath = getAbsolutePath(*Cmd).str().str();
+
+    if (!Loader.load(AbsPath, IndexStorage))
+      NeedsReIndexing.push_back({std::move(*Cmd), IndexStorage});
   }
-  // Load shard information into background-index.
-  {
-    std::lock_guard<std::mutex> Lock(ShardVersionsMu);
-    // 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.
-    for (const ShardInfo &SI : IntermediateSymbols) {
-      auto SS =
-          SI.Shard->Symbols
+
+  auto Shards = std::move(Loader).takeShards();
+  for (auto &SI : Shards) {
+    if (!SI.Shard)
+      continue;
+    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;
-      auto RelS =
-          SI.Shard->Relations
-              ? llvm::make_unique<RelationSlab>(std::move(*SI.Shard->Relations))
-              : nullptr;
-      ShardVersion &SV = ShardVersions[SI.AbsolutePath];
-      SV.Digest = SI.Digest;
-      SV.HadErrors = SI.HadErrors;
 
-      IndexedSymbols.update(SI.AbsolutePath, std::move(SS), std::move(RS),
-                            std::move(RelS), SI.CountReferences);
-    }
-  }
-  if (!IntermediateSymbols.empty())
+  // FIXME: Should we call it for each TU?
+  if (!Shards.empty())
     Rebuilder.loadedTU();
 
-  return Dependencies;
-}
-
-// Goes over each changed file and loads them from index. Returns the list of
-// TUs that had out-of-date/no shards.
-std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
-BackgroundIndex::loadShards(std::vector<std::string> ChangedFiles) {
-  std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
-      NeedsReIndexing;
-  // Keeps track of the files that will be reindexed, to make sure we won't
-  // re-index same dependencies more than once. Keys are AbsolutePaths.
-  llvm::StringSet<> FilesToIndex;
-  // Keeps track of the loaded shards to make sure we don't perform redundant
-  // disk IO. Keys are absolute paths.
-  llvm::StringSet<> LoadedShards;
-  Rebuilder.startLoading();
-  for (const auto &File : ChangedFiles) {
-    auto Cmd = CDB.getCompileCommand(File);
-    if (!Cmd)
-      continue;
-
-    std::string ProjectRoot;
-    if (auto PI = CDB.getProjectInfo(File))
-      ProjectRoot = std::move(PI->SourceRoot);
-
-    BackgroundIndexStorage *IndexStorage = IndexStorageFactory(ProjectRoot);
-    auto Dependencies = loadShard(*Cmd, IndexStorage, LoadedShards);
-    for (const auto &Dependency : Dependencies) {
-      if (!Dependency.NeedsReIndexing || FilesToIndex.count(Dependency.Path))
-        continue;
-      // FIXME: Currently, we simply schedule indexing on a TU whenever any of
-      // its dependencies needs re-indexing. We might do it smarter by figuring
-      // out a minimal set of TUs that will cover all the stale dependencies.
-      vlog("Enqueueing TU {0} because its dependency {1} needs re-indexing.",
-           Cmd->Filename, Dependency.Path);
-      NeedsReIndexing.push_back({std::move(*Cmd), IndexStorage});
-      // Mark all of this TU's dependencies as to-be-indexed so that we won't
-      // try to re-index those.
-      for (const auto &Dependency : Dependencies)
-        FilesToIndex.insert(Dependency.Path);
-      break;
-    }
-  }
   Rebuilder.doneLoading();
   return NeedsReIndexing;
 }
 
 } // namespace clangd
 } // namespace clang

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

+//===--- Background.h - Build an index in a background thread ----*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_INDEX_LOADER_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_INDEX_LOADER_H
+
+#include "Path.h"
+#include "index/Background.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/VirtualFileSystem.h"
+#include <memory>
+#include <vector>
+
+namespace clang {
+namespace clangd {
+
+struct ShardInfo {
+  Path AbsolutePath;
+  FileDigest Digest = {};
+  bool CountReferences = false;
+  bool HadErrors = false;
+  std::unique_ptr<IndexFileIn> Shard;
+};
+
+class BackgroundIndexLoader {
+public:
+  BackgroundIndexLoader(llvm::vfs::FileSystem *FS) : FS(FS) {}
+  /// Loads all shards for the TU \p MainFile from \p Storage. Returns false if
+  /// any shards that are first seen for this TU are stale.
+  bool load(PathRef MainFile, BackgroundIndexStorage *Storage);
+
+  /// Consumes the loader and returns all shards.
+  std::vector<ShardInfo> takeShards() &&;
+
+private:
+  struct CachedShard {
+    ShardInfo SI;
+    // Strings points to ShardInfo.
+    std::vector<Path> Edges;
+    bool IsStale = false;
+  };
+  std::pair<const CachedShard &, /*WasCached=*/bool>
+  loadShard(PathRef ShardIdentifier, BackgroundIndexStorage *Storage);
+
+  // Cache for Storage lookups.
+  llvm::StringMap<CachedShard> LoadedShards;
+
+  llvm::vfs::FileSystem *FS;
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif

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

+//===-- BackgroundIndexLoader.cpp - ---------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "index/BackgroundIndexLoader.h"
+#include "Logger.h"
+#include "Path.h"
+#include "index/Background.h"
+#include "llvm/ADT/DenseSet.h"
+#include <vector>
+
+namespace clang {
+namespace clangd {
+namespace {
+
+llvm::Optional<Path> uriToAbsolutePath(llvm::StringRef URI, PathRef HintPath) {
+  auto U = URI::parse(URI);
+  if (!U)
+    return llvm::None;
+  auto AbsolutePath = URI::resolve(*U, HintPath);
+  if (!AbsolutePath)
+    return llvm::None;
+  return *AbsolutePath;
+}
+
+bool hasChanged(llvm::vfs::FileSystem *FS, const ShardInfo &SI) {
+  auto Buf = FS->getBufferForFile(SI.AbsolutePath);
+  if (!Buf) {
+    elog("Couldn't get buffer for file: {0}: {1}", SI.AbsolutePath,
+         Buf.getError().message());
+    // There is no point in indexing an unreadable file.
+    return false;
+  }
+  return digest(Buf->get()->getBuffer()) != SI.Digest;
+}
+
+} // namespace
+
+std::pair<const BackgroundIndexLoader::CachedShard &, /*WasCached=*/bool>
+BackgroundIndexLoader::loadShard(PathRef ShardIdentifier,
+                                 BackgroundIndexStorage *Storage) {
+  auto It = LoadedShards.try_emplace(ShardIdentifier);
+  CachedShard &CS = It.first->getValue();
+  if (!It.second)
+    return {CS, true};
+
+  ShardInfo &SI = CS.SI;
+  SI.AbsolutePath = ShardIdentifier;
+  auto Shard = Storage->loadShard(ShardIdentifier);
+  if (!Shard || !Shard->Sources) {
+    vlog("Failed to load shard: {0}", ShardIdentifier);
+    SI.HadErrors = true;
+    CS.IsStale = true;
+    return {CS, false};
+  }
+
+  SI.Shard = std::move(Shard);
+  SI.AbsolutePath = ShardIdentifier;
+  for (const auto &It : *SI.Shard->Sources) {
+    auto AbsPath = uriToAbsolutePath(It.getKey(), ShardIdentifier);
+    if (!AbsPath || *AbsPath != ShardIdentifier) {
+      CS.Edges.push_back(*AbsPath);
+      continue;
+    }
+
+    const IncludeGraphNode &IGN = It.getValue();
+    SI.Digest = IGN.Digest;
+    SI.CountReferences = IGN.Flags & IncludeGraphNode::SourceFlag::IsTU;
+    SI.HadErrors = IGN.Flags & IncludeGraphNode::SourceFlag::HadErrors;
+  }
+  assert(SI.Digest != FileDigest{{0}} && "Digest is empty?");
+  CS.IsStale = hasChanged(FS, SI);
+  return {CS, false};
+}
+
+bool BackgroundIndexLoader::load(PathRef MainFile,
+                                 BackgroundIndexStorage *Storage) {
+  bool IsStale = false;
+  // Following containers points to strings inside LoadedShards.
+  llvm::DenseSet<PathRef> InQueue;
+  std::queue<PathRef> ToVisit;
+  auto LoadResult = loadShard(MainFile, Storage);
+
+  const CachedShard &MainShard = LoadResult.first;
+  if (!LoadResult.second && MainShard.IsStale)
+    IsStale = true;
+
+  for (PathRef Edge : MainShard.Edges) {
+    InQueue.insert(Edge);
+    ToVisit.push(Edge);
+  }
+  while (!ToVisit.empty()) {
+    PathRef ShardIdentifier = ToVisit.front();
+    ToVisit.pop();
+
+    auto LoadResult = loadShard(ShardIdentifier, Storage);
+    const CachedShard &CS = LoadResult.first;
+    // Report the file as stale if it is the first time we see it.
+    // FIXME: We should still update staleness for main file, if it had errors.
+    if (!LoadResult.second && CS.IsStale)
+      IsStale = true;
+    for (PathRef Edge : CS.Edges) {
+      if (InQueue.insert(Edge).second)
+        ToVisit.push(Edge);
+    }
+  }
+
+  return !IsStale;
+}
+
+std::vector<ShardInfo> BackgroundIndexLoader::takeShards() && {
+  std::vector<ShardInfo> Shards;
+  Shards.reserve(LoadedShards.size());
+  for (auto &It : LoadedShards)
+    Shards.push_back(std::move(It.getValue().SI));
+  return Shards;
+}
+
+} // namespace clangd
+} // namespace clang