[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/clangd/index/BackgroundIndexLoader.cpp

public:
   /// 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);
+  loadShard(PathRef StartSourceFile, PathRef DependentTU);
 
   /// Cache for Storage lookups.
   llvm::StringMap<LoadedShard> LoadedShards;
 
-  /// References are into the AbsolutePaths in LoadedShards.
-  llvm::DenseMap<PathRef, PathRef> FileToTU;
-
   BackgroundIndexStorage::Factory &IndexStorageFactory;
 };
 
 std::pair<const LoadedShard &, std::vector<Path>>
-BackgroundIndexLoader::loadShard(PathRef StartSourceFile) {
+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);
void BackgroundIndexLoader::load(PathRef MainFile) {
   std::queue<PathRef> ToVisit;
   InQueue.insert(MainFile);
   ToVisit.push(MainFile);
 
   while (!ToVisit.empty()) {
     PathRef SourceFile = ToVisit.front();
     ToVisit.pop();
 
-    auto ShardAndEdges = loadShard(SourceFile);
-    FileToTU[ShardAndEdges.first.AbsolutePath] = MainFile;
+    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) {
+  for (auto &It : LoadedShards)
     Result.push_back(std::move(It.getValue()));
-    LoadedShard &LS = Result.back();
-    auto TUsIt = FileToTU.find(LS.AbsolutePath);
-    assert(TUsIt != FileToTU.end() && "No TU registered for the shard");
-    Result.back().DependentTU = TUsIt->second;
-  }
   return Result;
 }
 } // namespace
 
 std::vector<LoadedShard>
 loadIndexShards(llvm::ArrayRef<Path> MainFiles,
                 BackgroundIndexStorage::Factory &IndexStorageFactory,
                 const GlobalCompilationDatabase &CDB) {