[clangd][BackgroundIndexLoader] Directly store DependentTU while loading shard

[kadircet]

Summary

We were deferring the population of DependentTU field in LoadedShard
until BackgroundIndexLoader was consumed. This actually triggers a use after
free since the shards FileToTU was pointing at could've been moved while
consuming the Loader.

clang-tools-extra/trunk/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/trunk/clangd/ClangdServer.cpp

if (this->Index != nullptr) {
       this->Index = Idx;
     }
   };
   if (Opts.StaticIndex)
     AddIndex(Opts.StaticIndex);
   if (Opts.BackgroundIndex) {
     BackgroundIdx = llvm::make_unique<BackgroundIndex>(
         Context::current().clone(), FSProvider, CDB,
-        BackgroundIndexStorage::createDiskBackedStorageFactory());
+        BackgroundIndexStorage::createDiskBackedStorageFactory(
+            [&CDB](llvm::StringRef File) { return CDB.getProjectInfo(File); }));
     AddIndex(BackgroundIdx.get());
   }
   if (DynamicIdx)
     AddIndex(DynamicIdx.get());
 }
 
 void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
                                WantDiagnostics WantDiags) {

clang-tools-extra/trunk/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"
public:
   virtual llvm::Error storeShard(llvm::StringRef ShardIdentifier,
                                  IndexFileOut Shard) const = 0;
 
   // Tries to load shard with given identifier, returns nullptr if shard
   // couldn't be loaded.
   virtual std::unique_ptr<IndexFileIn>
   loadShard(llvm::StringRef ShardIdentifier) const = 0;
 
-  // The factory provides storage for each CDB.
+  // The factory provides storage for each File.
   // It keeps ownership of the storage instances, and should manage caching
   // itself. Factory must be threadsafe and never returns nullptr.
-  using Factory =
-      llvm::unique_function<BackgroundIndexStorage *(llvm::StringRef)>;
+  using Factory = llvm::unique_function<BackgroundIndexStorage *(PathRef)>;
 
   // Creates an Index Storage that saves shards into disk. Index storage uses
-  // CDBDirectory + ".clangd/index/" as the folder to save shards.
-  static Factory createDiskBackedStorageFactory();
+  // CDBDirectory + ".clangd/index/" as the folder to save shards. CDBDirectory
+  // is the first directory containing a CDB in parent directories of a file, or
+  // user's home directory if none was found, e.g. standard library headers.
+  static Factory createDiskBackedStorageFactory(
+      std::function<llvm::Optional<ProjectInfo>(PathRef)> GetProjectInfo);
 };
 
 // A priority queue of tasks which can be run on (external) worker threads.
 class BackgroundQueue {
 public:
   /// A work item on the thread pool's queue.
   struct Task {
     explicit Task(std::function<void()> Run) : Run(std::move(Run)) {}
struct ShardVersion {
     bool HadErrors = false;
   };
 
   /// Given index results from a TU, only update symbols coming from files with
   /// different digests than \p ShardVersionsSnapshot. Also stores new index
   /// information on IndexStorage.
   void update(llvm::StringRef MainFile, IndexFileIn Index,
               const llvm::StringMap<ShardVersion> &ShardVersionsSnapshot,
-              BackgroundIndexStorage *IndexStorage, bool HadErrors);
+              bool HadErrors);
 
   // configuration
   const FileSystemProvider &FSProvider;
   const GlobalCompilationDatabase &CDB;
   Context BackgroundContext;
 
-  llvm::Error index(tooling::CompileCommand,
-                    BackgroundIndexStorage *IndexStorage);
+  llvm::Error index(tooling::CompileCommand);
 
   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.
-  std::vector<std::pair<tooling::CompileCommand, BackgroundIndexStorage *>>
-  loadShards(std::vector<std::string> ChangedFiles);
+  // Tries to load shards for the MainFiles and their dependencies.
+  std::vector<tooling::CompileCommand>
+  loadProject(std::vector<std::string> MainFiles);
 
   BackgroundQueue::Task
   changedFilesTask(const std::vector<std::string> &ChangedFiles);
-  BackgroundQueue::Task indexFileTask(tooling::CompileCommand Cmd,
-                                      BackgroundIndexStorage *Storage);
+  BackgroundQueue::Task indexFileTask(tooling::CompileCommand Cmd);
 
   // from lowest to highest priority
   enum QueuePriority {
     IndexFile,
     IndexBoostedFile,
     LoadShards,
   };
   BackgroundQueue Queue;
   AsyncTaskRunner ThreadPool;
   GlobalCompilationDatabase::CommandChanged::Subscription CommandsChanged;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/trunk/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 "clang/Driver/Types.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.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;
 }
+
+bool shardIsStale(const LoadedShard &LS, llvm::vfs::FileSystem *FS) {
+  auto Buf = FS->getBufferForFile(LS.AbsolutePath);
+  if (!Buf) {
+    elog("Background-index: Couldn't read {0} to validate stored index: {1}",
+         LS.AbsolutePath, Buf.getError().message());
+    // There is no point in indexing an unreadable file.
+    return false;
+  }
+  return digest(Buf->get()->getBuffer()) != LS.Digest;
+}
+
 } // 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)),

 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));
+    for (auto &Cmd : NeedsReIndexing)
+      Tasks.push_back(indexFileTask(std::move(Cmd)));
     Queue.append(std::move(Tasks));
   });
 
   T.QueuePri = LoadShards;
   T.ThreadPri = llvm::ThreadPriority::Default;
   return T;
 }
 
 static llvm::StringRef filenameWithoutExtension(llvm::StringRef Path) {
   Path = llvm::sys::path::filename(Path);
   return Path.drop_back(llvm::sys::path::extension(Path).size());
 }
 
 BackgroundQueue::Task
-BackgroundIndex::indexFileTask(tooling::CompileCommand Cmd,
-                               BackgroundIndexStorage *Storage) {
-  BackgroundQueue::Task T([this, Storage, Cmd] {
+BackgroundIndex::indexFileTask(tooling::CompileCommand Cmd) {
+  BackgroundQueue::Task T([this, Cmd] {
     // We can't use llvm::StringRef here since we are going to
     // move from Cmd during the call below.
     const std::string FileName = Cmd.Filename;
-    if (auto Error = index(std::move(Cmd), Storage))
+    if (auto Error = index(std::move(Cmd)))
       elog("Indexing {0} failed: {1}", FileName, std::move(Error));
   });
   T.QueuePri = IndexFile;
   T.Tag = filenameWithoutExtension(Cmd.Filename);
   return T;
 }
 
 void BackgroundIndex::boostRelated(llvm::StringRef Path) {
   namespace types = clang::driver::types;
   auto Type =
       types::lookupTypeForExtension(llvm::sys::path::extension(Path).substr(1));
   // is this a header?
   if (Type != types::TY_INVALID && types::onlyPrecompileType(Type))
     Queue.boost(filenameWithoutExtension(Path), IndexBoostedFile);
 }
 
 /// Given index results from a TU, only update symbols coming from files that
 /// are different or missing from than \p ShardVersionsSnapshot. Also stores new
 /// index information on IndexStorage.
 void BackgroundIndex::update(
     llvm::StringRef MainFile, IndexFileIn Index,
     const llvm::StringMap<ShardVersion> &ShardVersionsSnapshot,
-    BackgroundIndexStorage *IndexStorage, bool HadErrors) {
+    bool HadErrors) {
   // Partition symbols/references into files.
   struct File {
     llvm::DenseSet<const Symbol *> Symbols;
     llvm::DenseSet<const Ref *> Refs;
     llvm::DenseSet<const Relation *> Relations;
     FileDigest Digest;
   };
   llvm::StringMap<File> Files;
for (const auto &FileIt : Files) {
     auto RS = llvm::make_unique<RefSlab>(std::move(Refs).build());
     auto RelS = llvm::make_unique<RelationSlab>(std::move(Relations).build());
     auto IG = llvm::make_unique<IncludeGraph>(
         getSubGraph(URI::create(Path), Index.Sources.getValue()));
 
     // We need to store shards before updating the index, since the latter
     // consumes slabs.
     // FIXME: Also skip serializing the shard if it is already up-to-date.
-    if (IndexStorage) {
-      IndexFileOut Shard;
-      Shard.Symbols = SS.get();
-      Shard.Refs = RS.get();
-      Shard.Relations = RelS.get();
-      Shard.Sources = IG.get();
+    BackgroundIndexStorage *IndexStorage = IndexStorageFactory(Path);
 
-      // Only store command line hash for main files of the TU, since our
-      // current model keeps only one version of a header file.
-      if (Path == MainFile)
-        Shard.Cmd = Index.Cmd.getPointer();
 
-      if (auto Error = IndexStorage->storeShard(Path, Shard))
-        elog("Failed to write background-index shard for file {0}: {1}", Path,
-             std::move(Error));
-    }
 
     {
       std::lock_guard<std::mutex> Lock(ShardVersionsMu);
       auto Hash = FileIt.second.Digest;
       auto DigestIt = ShardVersions.try_emplace(Path);
       ShardVersion &SV = DigestIt.first->second;
       // Skip if file is already up to date, unless previous index was broken
       // and this one is not.
       if (!DigestIt.second && SV.Digest == Hash && SV.HadErrors && !HadErrors)
         continue;
       SV.Digest = Hash;
       SV.HadErrors = HadErrors;
 
       // 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), std::move(RelS),
                             Path == MainFile);
     }
   }
 }
 
-llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd,
-                                   BackgroundIndexStorage *IndexStorage) {
+llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
   trace::Span Tracer("BackgroundIndex");
   SPAN_ATTACH(Tracer, "file", Cmd.Filename);
   auto AbsolutePath = getAbsolutePath(Cmd);
 
   auto FS = FSProvider.getFileSystem();
   auto Buf = FS->getBufferForFile(AbsolutePath);
   if (!Buf)
     return llvm::errorCodeToError(Buf.getError());
llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
 
   bool HadErrors = Clang->hasDiagnostics() &&
                    Clang->getDiagnostics().hasUncompilableErrorOccurred();
   if (HadErrors) {
     log("Failed to compile {0}, index may be incomplete", AbsolutePath);
     for (auto &It : *Index.Sources)
       It.second.Flags |= IncludeGraphNode::SourceFlag::HadErrors;
   }
-  update(AbsolutePath, std::move(Index), ShardVersionsSnapshot, IndexStorage,
-         HadErrors);
+  update(AbsolutePath, std::move(Index), ShardVersionsSnapshot, 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;
-  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;
-    }
+// Restores shards for \p MainFiles from index storage. Then checks staleness of
+// those shards and returns a list of TUs that needs to be indexed to update
+// staleness.
+std::vector<tooling::CompileCommand>
+BackgroundIndex::loadProject(std::vector<std::string> MainFiles) {
+  std::vector<tooling::CompileCommand> NeedsReIndexing;
 
-    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;
-
-      // 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;
-    }
-  }
-  // Load shard information into background-index.
+  Rebuilder.startLoading();
+  // Load shards for all of the mainfiles.
+  const std::vector<LoadedShard> Result =
+      loadIndexShards(MainFiles, IndexStorageFactory, CDB);
+  size_t LoadedShards = 0;
   {
+    // Update in-memory state.
     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) {
+    for (auto &LS : Result) {
+      if (!LS.Shard)
+        continue;
       auto SS =
-          SI.Shard->Symbols
-              ? llvm::make_unique<SymbolSlab>(std::move(*SI.Shard->Symbols))
+          LS.Shard->Symbols
+              ? llvm::make_unique<SymbolSlab>(std::move(*LS.Shard->Symbols))
               : nullptr;
-      auto RS = SI.Shard->Refs
-                    ? llvm::make_unique<RefSlab>(std::move(*SI.Shard->Refs))
+      auto RS = LS.Shard->Refs
+                    ? llvm::make_unique<RefSlab>(std::move(*LS.Shard->Refs))
                     : nullptr;
       auto RelS =
-          SI.Shard->Relations
-              ? llvm::make_unique<RelationSlab>(std::move(*SI.Shard->Relations))
+          LS.Shard->Relations
+              ? llvm::make_unique<RelationSlab>(std::move(*LS.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())
-    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;
+      ShardVersion &SV = ShardVersions[LS.AbsolutePath];
+      SV.Digest = LS.Digest;
+      SV.HadErrors = LS.HadErrors;
+      ++LoadedShards;
 
-    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;
-    }
+      IndexedSymbols.update(LS.AbsolutePath, std::move(SS), std::move(RS),
+                            std::move(RelS), LS.CountReferences);
+    }
+  }
+  Rebuilder.loadedShard(LoadedShards);
+  Rebuilder.doneLoading();
+
+  auto FS = FSProvider.getFileSystem();
+  llvm::DenseSet<PathRef> TUsToIndex;
+  // We'll accept data from stale shards, but ensure the files get reindexed
+  // soon.
+  for (auto &LS : Result) {
+    if (!shardIsStale(LS, FS.get()))
+    PathRef TUForFile = LS.DependentTU;
+    assert(!TUForFile.empty() && "File without a TU!");
+
+    // FIXME: Try looking at other TUs if no compile commands are available
+    // for this TU, i.e TU was deleted after we performed indexing.
+    TUsToIndex.insert(TUForFile);
+
+  for (PathRef TU : TUsToIndex) {
+    auto Cmd = CDB.getCompileCommand(TU);
+    if (!Cmd)
+      continue;
+    NeedsReIndexing.emplace_back(std::move(*Cmd));
   }
-  Rebuilder.doneLoading();
+
   return NeedsReIndexing;
 }
 
 } // namespace clangd
 } // namespace clang

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

+//===--- BackgroundIndexLoader.h - Load shards from index storage-*- 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/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/VirtualFileSystem.h"
+#include <memory>
+#include <vector>
+
+namespace clang {
+namespace clangd {
+
+/// Represents a shard loaded from storage, stores contents in \p Shard and
+/// metadata about the source file that generated this shard.
+struct LoadedShard {
+  /// Path of the source file that produced this shard.
+  Path AbsolutePath;
+  /// Digest of the source file contents that produced this shard.
+  FileDigest Digest = {};
+  /// Whether the RefSlab in Shard should be used for updating symbol reference
+  /// counts when building an index.
+  bool CountReferences = false;
+  /// Whether the indexing action producing that shard had errors.
+  bool HadErrors = false;
+  /// Path to a TU that is depending on this shard.
+  Path DependentTU;
+  /// Will be nullptr when index storage couldn't provide a valid shard for
+  /// AbsolutePath.
+  std::unique_ptr<IndexFileIn> Shard;
+};
+
+/// Loads all shards for the TU \p MainFile from \p Storage.
+std::vector<LoadedShard>
+loadIndexShards(llvm::ArrayRef<Path> MainFiles,
+                BackgroundIndexStorage::Factory &IndexStorageFactory,
+                const GlobalCompilationDatabase &CDB);
+
+} // namespace clangd
+} // namespace clang
+
+#endif

clang-tools-extra/trunk/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 "GlobalCompilationDatabase.h"
+#include "Logger.h"
+#include "Path.h"
+#include "index/Background.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Path.h"
+#include <string>
+#include <utility>
+#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;
+}
+
+/// A helper class to cache BackgroundIndexStorage operations and keep the
+/// inverse dependency mapping.
+class BackgroundIndexLoader {
+public:
+  BackgroundIndexLoader(BackgroundIndexStorage::Factory &IndexStorageFactory)
+      : IndexStorageFactory(IndexStorageFactory) {}
+  /// Load the shards for \p MainFile and all of its dependencies.
+  void load(PathRef MainFile);
+
+  /// Consumes the loader and returns all shards.
+  std::vector<LoadedShard> takeResult() &&;
+
+private:
+  /// Returns the Shard for \p StartSourceFile from cache or loads it from \p
+  /// Storage. Also returns paths for dependencies of \p StartSourceFile if it
+  /// wasn't cached yet.
+  std::pair<const LoadedShard &, std::vector<Path>>
+  loadShard(PathRef StartSourceFile, PathRef DependentTU);
+
+  /// Cache for Storage lookups.
+  llvm::StringMap<LoadedShard> LoadedShards;
+
+  BackgroundIndexStorage::Factory &IndexStorageFactory;
+};
+
+std::pair<const LoadedShard &, std::vector<Path>>
+BackgroundIndexLoader::loadShard(PathRef StartSourceFile, PathRef DependentTU) {
+  auto It = LoadedShards.try_emplace(StartSourceFile);
+  LoadedShard &LS = It.first->getValue();
+  std::vector<Path> Edges = {};
+  // Return the cached shard.
+  if (!It.second)
+    return {LS, Edges};
+
+  LS.AbsolutePath = StartSourceFile.str();
+  LS.DependentTU = DependentTU;
+  BackgroundIndexStorage *Storage = IndexStorageFactory(LS.AbsolutePath);
+  auto Shard = Storage->loadShard(StartSourceFile);
+  if (!Shard || !Shard->Sources) {
+    vlog("Failed to load shard: {0}", StartSourceFile);
+    return {LS, Edges};
+  }
+
+  LS.Shard = std::move(Shard);
+  for (const auto &It : *LS.Shard->Sources) {
+    auto AbsPath = uriToAbsolutePath(It.getKey(), StartSourceFile);
+    if (!AbsPath)
+      continue;
+    // A shard contains only edges for non main-file sources.
+    if (*AbsPath != StartSourceFile) {
+      Edges.push_back(*AbsPath);
+      continue;
+    }
+
+    // Fill in shard metadata.
+    const IncludeGraphNode &IGN = It.getValue();
+    LS.Digest = IGN.Digest;
+    LS.CountReferences = IGN.Flags & IncludeGraphNode::SourceFlag::IsTU;
+    LS.HadErrors = IGN.Flags & IncludeGraphNode::SourceFlag::HadErrors;
+  }
+  assert(LS.Digest != FileDigest{{0}} && "Digest is empty?");
+  return {LS, Edges};
+}
+
+void BackgroundIndexLoader::load(PathRef MainFile) {
+  llvm::StringSet<> InQueue;
+  // Following containers points to strings inside InQueue.
+  std::queue<PathRef> ToVisit;
+  InQueue.insert(MainFile);
+  ToVisit.push(MainFile);
+
+  while (!ToVisit.empty()) {
+    PathRef SourceFile = ToVisit.front();
+    ToVisit.pop();
+
+    auto ShardAndEdges = loadShard(SourceFile, MainFile);
+    for (PathRef Edge : ShardAndEdges.second) {
+      auto It = InQueue.insert(Edge);
+      if (It.second)
+        ToVisit.push(It.first->getKey());
+    }
+  }
+}
+
+std::vector<LoadedShard> BackgroundIndexLoader::takeResult() && {
+  std::vector<LoadedShard> Result;
+  Result.reserve(LoadedShards.size());
+  for (auto &It : LoadedShards)
+    Result.push_back(std::move(It.getValue()));
+  return Result;
+}
+} // namespace
+
+std::vector<LoadedShard>
+loadIndexShards(llvm::ArrayRef<Path> MainFiles,
+                BackgroundIndexStorage::Factory &IndexStorageFactory,
+                const GlobalCompilationDatabase &CDB) {
+  BackgroundIndexLoader Loader(IndexStorageFactory);
+  for (llvm::StringRef MainFile : MainFiles) {
+    assert(llvm::sys::path::is_absolute(MainFile));
+    Loader.load(MainFile);
+  }
+  return std::move(Loader).takeResult();
+}
+
+} // namespace clangd
+} // namespace clang

clang-tools-extra/trunk/clangd/index/BackgroundIndexStorage.cpp

 //== BackgroundIndexStorage.cpp - Provide caching support to BackgroundIndex ==/
 //
 // 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 "GlobalCompilationDatabase.h"
 #include "Logger.h"
+#include "Path.h"
 #include "index/Background.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
+#include <functional>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 std::string getShardPathFromFilePath(llvm::StringRef ShardRoot,
                                      llvm::StringRef FilePath) {
   llvm::SmallString<128> ShardRootSS(ShardRoot);
vlog("Couldn't find project for {0}, indexing in-memory only",
          ShardIdentifier);
     return llvm::Error::success();
   }
 };
 
 // Creates and owns IndexStorages for multiple CDBs.
 class DiskBackedIndexStorageManager {
 public:
-  DiskBackedIndexStorageManager()
-      : IndexStorageMapMu(llvm::make_unique<std::mutex>()) {}
+  DiskBackedIndexStorageManager(
+      std::function<llvm::Optional<ProjectInfo>(PathRef)> GetProjectInfo)
+      : IndexStorageMapMu(llvm::make_unique<std::mutex>()),
+        GetProjectInfo(std::move(GetProjectInfo)) {
+    llvm::SmallString<128> HomeDir;
+    llvm::sys::path::home_directory(HomeDir);
+    this->HomeDir = HomeDir.str().str();
+  }
 
-  // Creates or fetches to storage from cache for the specified CDB.
-  BackgroundIndexStorage *operator()(llvm::StringRef CDBDirectory) {
+  // Creates or fetches to storage from cache for the specified project.
+  BackgroundIndexStorage *operator()(PathRef File) {
     std::lock_guard<std::mutex> Lock(*IndexStorageMapMu);
+    Path CDBDirectory = HomeDir;
+    if (auto PI = GetProjectInfo(File))
+      CDBDirectory = PI->SourceRoot;
     auto &IndexStorage = IndexStorageMap[CDBDirectory];
     if (!IndexStorage)
       IndexStorage = create(CDBDirectory);
     return IndexStorage.get();
   }
 
 private:
-  std::unique_ptr<BackgroundIndexStorage> create(llvm::StringRef CDBDirectory) {
-    if (CDBDirectory.empty())
+  std::unique_ptr<BackgroundIndexStorage> create(PathRef CDBDirectory) {
+    if (CDBDirectory.empty()) {
+      elog("Tried to create storage for empty directory!");
       return llvm::make_unique<NullStorage>();
+    }
     return llvm::make_unique<DiskBackedIndexStorage>(CDBDirectory);
   }
 
+  Path HomeDir;
+
   llvm::StringMap<std::unique_ptr<BackgroundIndexStorage>> IndexStorageMap;
   std::unique_ptr<std::mutex> IndexStorageMapMu;
+
+  std::function<llvm::Optional<ProjectInfo>(PathRef)> GetProjectInfo;
 };
 
 } // namespace
 
 BackgroundIndexStorage::Factory
-BackgroundIndexStorage::createDiskBackedStorageFactory() {
-  return DiskBackedIndexStorageManager();
+BackgroundIndexStorage::createDiskBackedStorageFactory(
+    std::function<llvm::Optional<ProjectInfo>(PathRef)> GetProjectInfo) {
+  return DiskBackedIndexStorageManager(std::move(GetProjectInfo));
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/trunk/clangd/index/BackgroundRebuild.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_INDEX_REBUILD_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_BACKGROUND_INDEX_REBUILD_H
 
 #include "index/FileIndex.h"
 #include "index/Index.h"
 #include "llvm/Support/Threading.h"
+#include <cstddef>
 
 namespace clang {
 namespace clangd {
 
 // The BackgroundIndexRebuilder builds the serving data structures periodically
 // in response to events in the background indexer. The goal is to ensure the
 // served data stays fairly fresh, without wasting lots of CPU rebuilding it
 // often.
public:
   // Called to indicate that all worker threads are idle.
   // May reindex, if the index is not up to date.
   void idle();
   // Called to indicate we're going to load a batch of shards from disk.
   // startLoading() and doneLoading() must be paired, but multiple loading
   // sessions may happen concurrently.
   void startLoading();
   // Called to indicate some shards were actually loaded from disk.
-  void loadedTU();
+  void loadedShard(size_t ShardCount);
   // Called to indicate we're finished loading shards from disk.
   // May rebuild (if any were loaded).
   void doneLoading();
 
   // Ensures we won't start any more rebuilds.
   void shutdown();
 
   // Thresholds for rebuilding as TUs get indexed.
private:
   // Index builds are versioned. ActiveVersion chases StartedVersion.
   unsigned StartedVersion = 0;
   unsigned ActiveVersion = 0;
   // How many TUs have we indexed so far since startup?
   unsigned IndexedTUs = 0;
   unsigned IndexedTUsAtLastRebuild = 0;
   // Are we loading shards? May be multiple concurrent sessions.
   unsigned Loading = 0;
-  unsigned LoadedTUs; // In the current loading session.
+  unsigned LoadedShards; // In the current loading session.
 
   SwapIndex *Target;
   FileSymbols *Source;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/trunk/clangd/index/BackgroundRebuild.cpp

maybeRebuild("when background indexer is idle", [this] {
     // (even if currently rebuilding! this ensures eventual completeness)
     return IndexedTUs > IndexedTUsAtLastRebuild;
   });
 }
 
 void BackgroundIndexRebuilder::startLoading() {
   std::lock_guard<std::mutex> Lock(Mu);
   if (!Loading)
-    LoadedTUs = 0;
+    LoadedShards = 0;
   ++Loading;
 }
-void BackgroundIndexRebuilder::loadedTU() {
+void BackgroundIndexRebuilder::loadedShard(size_t ShardCount) {
   std::lock_guard<std::mutex> Lock(Mu);
   assert(Loading);
-  ++LoadedTUs;
+  LoadedShards += ShardCount;
 }
 void BackgroundIndexRebuilder::doneLoading() {
   maybeRebuild("after loading index from disk", [this] {
     assert(Loading);
     --Loading;
     if (Loading)    // was loading multiple batches concurrently
       return false; // rebuild once the last batch is done.
     // Rebuild if we loaded any shards, or if we stopped an indexedTU rebuild.
-    return LoadedTUs > 0 || enoughTUsToRebuild();
+    return LoadedShards > 0 || enoughTUsToRebuild();
   });
 }
 
 void BackgroundIndexRebuilder::shutdown() {
   std::lock_guard<std::mutex> Lock(Mu);
   ShouldStop = true;
 }
 

clang-tools-extra/trunk/clangd/test/Inputs/background-index/definition.jsonrpc

 {
   "jsonrpc": "2.0",
   "method": "textDocument/didOpen",
   "params": {
     "textDocument": {
       "uri": "file://DIRECTORY/bar.cpp",
       "languageId": "cpp",
       "version": 1,
-      "text": "#include \"foo.h\"\nint main(){\nreturn foo();\n}"
+      "text": "#include \"sub_dir/foo.h\"\nint main(){\nreturn foo();\n}"
     }
   }
 }
 ---
 {
   "jsonrpc": "2.0",
   "id": 1,
   "method": "sync",

clang-tools-extra/trunk/clangd/test/Inputs/background-index/foo.h

-#ifndef FOO_H
-#define FOO_H
-int foo();
-#endif

clang-tools-extra/trunk/clangd/test/Inputs/background-index/foo.cpp

-#include "foo.h"
+#include "sub_dir/foo.h"
 int foo() { return 42; }

clang-tools-extra/trunk/clangd/test/Inputs/background-index/sub_dir/foo.h

+#ifndef FOO_H
+#define FOO_H
+int foo();
+#endif

clang-tools-extra/trunk/clangd/test/background-index.test

 # We need to splice paths into file:// URIs for this test.
 # UNSUPPORTED: win32
 
 # Use a copy of inputs, as we'll mutate it (as will the background index).
 # RUN: rm -rf %t
 # RUN: cp -r %S/Inputs/background-index %t
 # Need to embed the correct temp path in the actual JSON-RPC requests.
-# RUN: sed -i -e "s|DIRECTORY|%t|" %t/*
+# RUN: sed -i -e "s|DIRECTORY|%t|" %t/definition.jsonrpc
+# RUN: sed -i -e "s|DIRECTORY|%t|" %t/compile_commands.json
 
 # We're editing bar.cpp, which includes foo.h.
 # foo() is declared in foo.h and defined in foo.cpp.
 # The background index should allow us to go-to-definition on foo().
 # RUN: clangd -background-index -lit-test < %t/definition.jsonrpc | FileCheck %t/definition.jsonrpc
 
 # Test that the index is writing files in the expected location.
 # RUN: ls %t/.clangd/index/foo.cpp.*.idx
+# RUN: ls %t/sub_dir/.clangd/index/foo.h.*.idx
 
 # Test the index is read from disk: delete code and restart clangd.
 # RUN: rm %t/foo.cpp
 # RUN: clangd -background-index -lit-test < %t/definition.jsonrpc | FileCheck %t/definition.jsonrpc

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

TEST_F(BackgroundIndexTest, ShardStorageTest) {
   }
   EXPECT_EQ(CacheHits, 0U);
   EXPECT_EQ(Storage.size(), 2U);
 
   {
     OverlayCDB CDB(/*Base=*/nullptr);
     BackgroundIndex Idx(Context::empty(), FS, CDB,
                         [&](llvm::StringRef) { return &MSS; });
-    CDB.setCompileCommand(testPath("root"), Cmd);
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
     ASSERT_TRUE(Idx.blockUntilIdleForTest());
   }
   EXPECT_EQ(CacheHits, 2U); // Check both A.cc and A.h loaded from cache.
   EXPECT_EQ(Storage.size(), 2U);
 
   auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
   EXPECT_NE(ShardHeader, nullptr);
   EXPECT_THAT(
FS.Files[testPath("root/A.h")] = R"cpp(
       void f_b();
       class A_CC {};
       class A_CCnew {};
       )cpp";
   {
     OverlayCDB CDB(/*Base=*/nullptr);
     BackgroundIndex Idx(Context::empty(), FS, CDB,
                         [&](llvm::StringRef) { return &MSS; });
-    CDB.setCompileCommand(testPath("root"), Cmd);
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
     ASSERT_TRUE(Idx.blockUntilIdleForTest());
   }
   EXPECT_EQ(CacheHits, 2U); // Check both A.cc and A.h loaded from cache.
 
   // Check if the new symbol has arrived.
   auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
   EXPECT_NE(ShardHeader, nullptr);
   EXPECT_THAT(*ShardHeader->Symbols, Contains(Named("A_CCnew")));
 
   // Change source.
   FS.Files[testPath("root/A.cc")] =
       "#include \"A.h\"\nvoid g() { (void)common; }\nvoid f_b() {}";
   {
     CacheHits = 0;
     OverlayCDB CDB(/*Base=*/nullptr);
     BackgroundIndex Idx(Context::empty(), FS, CDB,
                         [&](llvm::StringRef) { return &MSS; });
-    CDB.setCompileCommand(testPath("root"), Cmd);
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
     ASSERT_TRUE(Idx.blockUntilIdleForTest());
   }
   EXPECT_EQ(CacheHits, 2U); // Check both A.cc and A.h loaded from cache.
 
   // Check if the new symbol has arrived.
   ShardHeader = MSS.loadShard(testPath("root/A.h"));
   EXPECT_NE(ShardHeader, nullptr);
   EXPECT_THAT(*ShardHeader->Symbols, Contains(Named("A_CCnew")));
TEST_F(BackgroundIndexRebuilderTest, IndexingTUs) {
   EXPECT_TRUE(checkRebuild([&] { Rebuilder.indexedTU(); }));
   for (unsigned I = 0; I < Rebuilder.TUsBeforeRebuild - 1; ++I)
     EXPECT_FALSE(checkRebuild([&] { Rebuilder.indexedTU(); }));
   EXPECT_TRUE(checkRebuild([&] { Rebuilder.indexedTU(); }));
 }
 
 TEST_F(BackgroundIndexRebuilderTest, LoadingShards) {
   Rebuilder.startLoading();
-  Rebuilder.loadedTU();
-  Rebuilder.loadedTU();
+  Rebuilder.loadedShard(10);
+  Rebuilder.loadedShard(20);
   EXPECT_TRUE(checkRebuild([&] { Rebuilder.doneLoading(); }));
 
   // No rebuild for no shards.
   Rebuilder.startLoading();
   EXPECT_FALSE(checkRebuild([&] { Rebuilder.doneLoading(); }));
 
   // Loads can overlap.
   Rebuilder.startLoading();
-  Rebuilder.loadedTU();
+  Rebuilder.loadedShard(1);
   Rebuilder.startLoading();
-  Rebuilder.loadedTU();
+  Rebuilder.loadedShard(1);
   EXPECT_FALSE(checkRebuild([&] { Rebuilder.doneLoading(); }));
-  Rebuilder.loadedTU();
+  Rebuilder.loadedShard(1);
   EXPECT_TRUE(checkRebuild([&] { Rebuilder.doneLoading(); }));
 
   // No rebuilding for indexed files while loading.
   Rebuilder.startLoading();
   for (unsigned I = 0; I < 3 * Rebuilder.TUsBeforeRebuild; ++I)
     EXPECT_FALSE(checkRebuild([&] { Rebuilder.indexedTU(); }));
   // But they get indexed when we're done, even if no shards were loaded.
   EXPECT_TRUE(checkRebuild([&] { Rebuilder.doneLoading(); }));