[clangd] Put direct headers into srcs section.

[kadircet]

Summary

Currently, there's no way of knowing about header files
using compilation database, since it doesn't contain header files as entries.

Using this information, restoring from cache using compile commands becomes
possible instead of doing directory traversal. Also, we can issue indexing
actions for out-of-date headers even if source files depending on them haven't
changed.

[kadircet]

• Use include depth to get includes for all files, rather than just main file.

[sammccall]

Going to leave awkward comments suggesting we expand the scope a bit.
Full disclosure: file dependency information is something that's come up as useful in lots of contexts.

clangd/index/Serialization.cpp

501

This seems like unneccesary coupling between layers. On the contrary, I think it's going to be really useful to have file and dependency information in the index data and in the in-memory index.
(e.g. for #include completion, more accurately guessing compile commands for headers, deciding which CC file corresponds to a .h file and vice versa...)

clangd/index/Serialization.h

54

This seems a little off to me conceptually: an index contains symbols that might be from many files. (Also true of Digest, which I missed).

I'd suggest a richer structure which can represent one or many files:

std::unordered_map</*uri*/string, SourceFile>; // keys are URIs
struct SourceFile {
  bool IsTU;
  StringRef URI; // points back into key
  FileDigest Digest;
  vector<StringRef> Headers; // points back into keys
  // maybe later: compile command
}

(though probably wrap the map into a class to ensure the string ownership doesn't get messed up)

For auto-index shards we could include full SourceFile info for the main file (with IsTU = true, and headers) and each of the headers would have a skeletal entry with IsTU = false and no headers stored.

But clangd-indexer could spit out an index that contains full information for all the transitive include of all indexed TUs.

[kadircet]

• Change srcs to store multiple source file information rathet than one.

[kadircet]

• rebased and changed URI convertion patterns, since it changed in rL347467

[kadircet]

• Rebase to get changes in background index.

[sammccall]

structure and serialization stuff looks great!
still some questions around the auto-index logic and organization of the extraction code, let's chat offline. May make sense to split.

clangd/index/Background.cpp

329

nit: lowercase

341

It looks like you're re-reading the files here. this won't reuse in-memory copies of the file, and won't respect overridden contents outside the VFS (ugh). The SourceManager would be better I think.

351

I don't think this is going to work:

• you're outputting the non-transitive dependencies from the TU, but I think you want the *transitive* ones, which should then be partitioned by file (same as symbols). Otherwise the headers from the shard for foo.h aren't going to be right, are they?
• I don't think the current IncludeStructure includes all the info you want for transitive includes. e.g. for include edges A->B where there is a shorter path from Main->B, we don't record anything.
• I think we will want the include information for other forms of index too. I think this means the logic should be next to symbolcollector, or in Headers.h, where it can be shared

clangd/index/Serialization.cpp

400

these string reassignments deserve a comment

clangd/index/Serialization.h

41

Does this belong here, vs in Index, or SymbolCollector, or so?

53

I think the StringMap could also use a typedef near SourceFile - that's where we could mention that FileURI and DirectIncludes point back at the keys.

[sammccall]

To summarize offline discussion: I think we need to split this patch up, as the scope is growing. (In a really useful direction I think!)

The first patch should add the include graph data structure. (I think the right place for it is probably Headers.h, as we want to reuse it). It can also include the code to serialize/deserialize it in index files, which is in good shape.

Then we have at least 3 independent dependent patches :-)

• populate the transitive include graph structure from a PPCallbacks. Probably hook it up to the static index FrontendAction too, since that's simple. This can be tested directly.
• write the include graph appropriately in auto-index. This isn't quite trivial, as it needs to be partitioned by file.
• read shards in auto-index. This could be all-at-once, or a simple version that didn't follow #include edges first, and then modify to consume the include graph.

This is a bunch more work than originally envisioned, but having dependency graph info in the index enables/improves a bunch of features. Making this reusable across indexes seems worth it.

[kadircet]

• Only keep the build graph structures and {de,}serialization logic.
• Rename a few structures, move to more relavant places.

[sammccall]

Only really significant thing is I think the name "build graph" is confusing, and we should specifically mention includes.
Rest is just nits to take or leave...

clangd/Headers.h

52 ↗

(On Diff #175469)

May be worth explicitly noting:
Does not own the strings it references (BuildGraphInclusions is self-contained)

53 ↗

(On Diff #175469)

I think we should be more specific and call this the "include graph".
Build graph can suggest other build-system related things (bazel and ninja have a very formal concept of a build graph, other systems have less formalized but may still use the term)

54 ↗

(On Diff #175469)

this is worth a comment (e.g. A "main file" as opposed to a header.)

55 ↗

(On Diff #175469)

This is slightly ambiguous: URI of the file, or URI with the file scheme?
I think you mean the former.
I think URI is actually the best name for this field, even though it's taken by the class. As a member of a struct, there's not really any ambiguous cases.

57 ↗

(On Diff #175469)

nit: I don't think "File" actually adds anything here, the whole struct is scoped to a file.
I'd add a comment that this is only direct inclusions, if you choose not to mention it in the name.
(The comment seems more useful here than it does on the struct)

clangd/SourceCode.h

26 ↗

(On Diff #175469)

Probably worth a comment, especially since there are no functions here using it.

[kadircet]

• Address comments.
• Move digest and digestFile into SourceCode.h

clangd/index/Background.h

public:
 
   // Wait until the queue is empty, to allow deterministic testing.
   void blockUntilIdleForTest();
 
   using FileDigest = decltype(llvm::SHA1::hash({}));
 
 private:
   /// Given index results from a TU, only update files in \p FilesToUpdate.
-  void update(llvm::StringRef MainFile, SymbolSlab Symbols, RefSlab Refs,
+  /// Also stores new index information on IndexStorage.
+  void update(llvm::StringRef MainFile, IndexFileIn Index,
               const llvm::StringMap<FileDigest> &FilesToUpdate,
               BackgroundIndexStorage *IndexStorage);
 
   // configuration
   std::string ResourceDir;
   const FileSystemProvider &FSProvider;
   Context BackgroundContext;
 

clangd/index/Background.cpp

 #include "index/Serialization.h"
 #include "index/SymbolCollector.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/SHA1.h"
+#include "llvm/Support/ScopedPrinter.h"
 
 #include <memory>
 #include <queue>
 #include <random>
 #include <string>
 
 using namespace llvm;
 namespace clang {
llvm::StringRef resolve(llvm::StringRef FileURI) {
     return I.first->second;
   }
 
 private:
   std::string HintPath;
   llvm::StringMap<std::string> URIToPathCache;
 };
 
+llvm::StringMap<IndexFileIn::SourceFile>
+FilterSources(llvm::StringRef Path,
+              const llvm::StringMap<IndexFileIn::SourceFile> &FullSources) {
+  llvm::StringMap<IndexFileIn::SourceFile> Filtered;
+
+  std::string FileURI = URI::create(Path).toString();
+  std::queue<llvm::StringRef> Dependencies;
+  Dependencies.push(FileURI);
+  while (!Dependencies.empty()) {
+    llvm::StringRef Dependency = Dependencies.front();
+    Dependencies.pop();
+
+    auto Entry = Filtered.try_emplace(Dependency).first;
+    auto &SF = Entry->getValue();
+    SF = FullSources.lookup(Entry->getKey());
+    SF.FileURI = Entry->getKey();
+
+    for (auto &Include : SF.DirectIncludes) {
+      auto I = Filtered.try_emplace(Include);
+      Include = I.first->getKey();
+      if (I.second) // Insert each dependency only once.
+        Dependencies.push(Include);
+    }
+  }
+
+  return Filtered;
+}
+
 /// Given index results from a TU, only update files in \p FilesToUpdate.
-void BackgroundIndex::update(StringRef MainFile, SymbolSlab Symbols,
-                             RefSlab Refs,
+void BackgroundIndex::update(StringRef MainFile, IndexFileIn Index,
                              const StringMap<FileDigest> &FilesToUpdate,
                              BackgroundIndexStorage *IndexStorage) {
   // Partition symbols/references into files.
   struct File {
     DenseSet<const Symbol *> Symbols;
     DenseSet<const Ref *> Refs;
   };
   StringMap<File> Files;
   URIToFileCache URICache(MainFile);
-  for (const auto &Sym : Symbols) {
+  for (const auto &Sym : *Index.Symbols) {
     if (Sym.CanonicalDeclaration) {
       auto DeclPath = URICache.resolve(Sym.CanonicalDeclaration.FileURI);
       if (FilesToUpdate.count(DeclPath) != 0)
         Files[DeclPath].Symbols.insert(&Sym);
     }
     // For symbols with different declaration and definition locations, we store
     // the full symbol in both the header file and the implementation file, so
     // that merging can tell the preferred symbols (from canonical headers) from
     // other symbols (e.g. forward declarations).
     if (Sym.Definition &&
         Sym.Definition.FileURI != Sym.CanonicalDeclaration.FileURI) {
       auto DefPath = URICache.resolve(Sym.Definition.FileURI);
       if (FilesToUpdate.count(DefPath) != 0)
         Files[DefPath].Symbols.insert(&Sym);
     }
   }
   DenseMap<const Ref *, SymbolID> RefToIDs;
-  for (const auto &SymRefs : Refs) {
+  for (const auto &SymRefs : *Index.Refs) {
     for (const auto &R : SymRefs.second) {
       auto Path = URICache.resolve(R.Location.FileURI);
       if (FilesToUpdate.count(Path) != 0) {
         auto &F = Files[Path];
         RefToIDs[&R] = SymRefs.first;
         F.Refs.insert(&R);
       }
     }
   }
 
   // Build and store new slabs for each updated file.
   for (const auto &F : Files) {
     StringRef Path = F.first();
     vlog("Update symbols in {0}", Path);
     SymbolSlab::Builder Syms;
     RefSlab::Builder Refs;
     for (const auto *S : F.second.Symbols)
       Syms.insert(*S);
     for (const auto *R : F.second.Refs)
       Refs.insert(RefToIDs[R], *R);
 
     auto SS = llvm::make_unique<SymbolSlab>(std::move(Syms).build());
     auto RS = llvm::make_unique<RefSlab>(std::move(Refs).build());
+    std::unique_ptr<llvm::StringMap<IndexFileIn::SourceFile>> Srcs;
 
     auto Hash = FilesToUpdate.lookup(Path);
     // We need to store shards before updating the index, since the latter
     // consumes slabs.
-    // FIXME: Store Hash in the Shard.
     if (IndexStorage) {
       IndexFileOut Shard;
       Shard.Symbols = SS.get();
       Shard.Refs = RS.get();
-      Shard.Digest = &Hash;
+      if (Index.Sources) {
+        Srcs = llvm::make_unique<llvm::StringMap<IndexFileIn::SourceFile>>(
+            FilterSources(Path, Index.Sources.getValue()));
+        Shard.Sources = Srcs.get();
+      }
+
       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(DigestsMu);
     // 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
return [&FileDigests, &FilesToUpdate](const SourceManager &SM, FileID FID) {
     if (D != FileDigests.end() && D->second == Digest)
       return false; // Skip files that haven't changed.
 
     FilesToUpdate[AbsPath] = *Digest;
     return true;
   };
 }
 
+llvm::StringMap<IndexFileIn::SourceFile>
+IncludeStructureToSources(llvm::StringRef MainFile,
+                          const IncludeStructure &Includes,
+                          llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS) {
+  llvm::StringMap<IndexFileIn::SourceFile> Sources;
+
+  std::queue<std::string> Dependencies;
+  Dependencies.push(MainFile);
+  while (!Dependencies.empty()) {
+    std::string Dependency = std::move(Dependencies.front());
+    Dependencies.pop();
+
+    std::string FileURI = URI::create(Dependency).toString();
+    auto Buf = FS->getBufferForFile(Dependency);
+    if (!Buf)
+      continue;
+    auto Digest = digest(Buf->get()->getBuffer());
+
+    auto Entry = Sources.try_emplace(FileURI).first;
+    auto &SF = Entry->getValue();
+    SF.IsTU = Dependency == MainFile;
+    SF.FileURI = Entry->getKey();
+    SF.Digest = std::move(Digest);
+    for (const auto &Include : Includes.includeDepth(Dependency)) {
+      if (Include.getValue() != 1) // Skip transitive includes.
+        continue;
+      FileURI = URI::create(Include.getKey()).toString();
+      auto I = Sources.try_emplace(FileURI);
+      SF.DirectIncludes.push_back(I.first->getKey());
+      if (I.second) // Insert each dependency only once.
+        Dependencies.push(Include.getKey());
+    }
+  }
+
+  return Sources;
+}
+
 Error BackgroundIndex::index(tooling::CompileCommand Cmd,
                              BackgroundIndexStorage *IndexStorage) {
   trace::Span Tracer("BackgroundIndex");
   SPAN_ATTACH(Tracer, "file", Cmd.Filename);
   SmallString<128> AbsolutePath;
   if (sys::path::is_absolute(Cmd.Filename)) {
     AbsolutePath = Cmd.Filename;
   } else {
Error BackgroundIndex::index(tooling::CompileCommand Cmd,
   // We could use CrashRecoveryContext to try to make indexing crashes nonfatal,
   // but the leaky "recovery" is pretty scary too in a long-running process.
   // If crashes are a real problem, maybe we should fork a child process.
 
   const FrontendInputFile &Input = Clang->getFrontendOpts().Inputs.front();
   if (!Action->BeginSourceFile(*Clang, Input))
     return createStringError(inconvertibleErrorCode(),
                              "BeginSourceFile() failed");
+
+  IncludeStructure Includes;
+  Clang->getPreprocessor().addPPCallbacks(
+      collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
+
   if (!Action->Execute())
     return createStringError(inconvertibleErrorCode(), "Execute() failed");
   Action->EndSourceFile();
 
+  llvm::StringMap<IndexFileIn::SourceFile> Sources =
+      IncludeStructureToSources(AbsolutePath, Includes, Inputs.FS);
+
   log("Indexed {0} ({1} symbols, {2} refs)", Inputs.CompileCommand.Filename,
       Symbols.size(), Refs.numRefs());
   SPAN_ATTACH(Tracer, "symbols", int(Symbols.size()));
   SPAN_ATTACH(Tracer, "refs", int(Refs.numRefs()));
-  update(AbsolutePath, std::move(Symbols), std::move(Refs), FilesToUpdate,
-         IndexStorage);
+  IndexFileIn Index;
+  Index.Symbols = std::move(Symbols);
+  Index.Refs = std::move(Refs);
+  Index.Sources = std::move(Sources);
+
+  update(AbsolutePath, std::move(Index), FilesToUpdate, IndexStorage);
   {
     // Make sure hash for the main file is always updated even if there is no
     // index data in it.
     std::lock_guard<std::mutex> Lock(DigestsMu);
     IndexedFileDigests[AbsolutePath] = Hash;
   }
 
   // FIXME: this should rebuild once-in-a-while, not after every file.

clangd/index/Serialization.h

 enum class IndexFileFormat {
   RIFF, // Versioned binary format, suitable for production use.
   YAML, // Human-readable format, suitable for experiments and debugging.
 };
 
 // Holds the contents of an index file that was read.
 struct IndexFileIn {
   using FileDigest = std::array<uint8_t, 20>;
+  // Contains information about one source file that participates in current
+  // index file.
+  struct SourceFile {
+    bool IsTU;
+    llvm::StringRef FileURI;
+    FileDigest Digest;
+    std::vector<llvm::StringRef> DirectIncludes;
+  };
+
   llvm::Optional<SymbolSlab> Symbols;
   llvm::Optional<RefSlab> Refs;
-  // Digest of the source file that generated the contents.
-  llvm::Optional<FileDigest> Digest;
+  // Keys are URIs of the source files.
+  llvm::Optional<llvm::StringMap<SourceFile>> Sources;
 };
 // Parse an index file. The input must be a RIFF or YAML file.
 llvm::Expected<IndexFileIn> readIndexFile(llvm::StringRef);
 
 // Specifies the contents of an index file to be written.
 struct IndexFileOut {
   const SymbolSlab *Symbols = nullptr;
   const RefSlab *Refs = nullptr;
-  // Digest of the source file that generated the contents.
-  const IndexFileIn::FileDigest *Digest = nullptr;
+  // Information about the source files participated in current index file.
+  // Keys are URIs of the source files.
+  const llvm::StringMap<IndexFileIn::SourceFile> *Sources = nullptr;
   // TODO: Support serializing Dex posting lists.
   IndexFileFormat Format = IndexFileFormat::RIFF;
 
   IndexFileOut() = default;
   IndexFileOut(const IndexFileIn &I)
       : Symbols(I.Symbols ? I.Symbols.getPointer() : nullptr),
         Refs(I.Refs ? I.Refs.getPointer() : nullptr) {}
 };

clangd/index/Serialization.cpp

SymbolLocation readLocation(Reader &Data, ArrayRef<StringRef> Strings) {
   Loc.FileURI = Data.consumeString(Strings).data();
   for (auto *Endpoint : {&Loc.Start, &Loc.End}) {
     Endpoint->setLine(Data.consumeVar());
     Endpoint->setColumn(Data.consumeVar());
   }
   return Loc;
 }
 
+IndexFileIn::SourceFile
+readSourceFile(Reader &Data, llvm::ArrayRef<llvm::StringRef> Strings) {
+  IndexFileIn::SourceFile SF;
+  SF.IsTU = Data.consume8();
+  SF.FileURI = Data.consumeString(Strings);
+  llvm::StringRef Digest = Data.consume(SF.Digest.size());
+  std::copy(Digest.bytes_begin(), Digest.bytes_end(), SF.Digest.begin());
+  SF.DirectIncludes.resize(Data.consumeVar());
+
+  for (llvm::StringRef& Include : SF.DirectIncludes)
+    Include = Data.consumeString(Strings);
+  return SF;
+}
+
+void writeSourceFile(const IndexFileIn::SourceFile &SF,
+                     const StringTableOut &Strings, raw_ostream &OS) {
+  OS.write(SF.IsTU);
+  writeVar(Strings.index(SF.FileURI), OS);
+  llvm::StringRef Hash(reinterpret_cast<const char *>(SF.Digest.data()),
+                       SF.Digest.size());
+  OS << Hash;
+  writeVar(SF.DirectIncludes.size(), OS);
+  for (llvm::StringRef Include : SF.DirectIncludes)
+    writeVar(Strings.index(Include), OS);
+}
+
 void writeSymbol(const Symbol &Sym, const StringTableOut &Strings,
                  raw_ostream &OS) {
   OS << Sym.ID.raw(); // TODO: once we start writing xrefs and posting lists,
                       // symbol IDs should probably be in a string table.
   OS.write(static_cast<uint8_t>(Sym.SymInfo.Kind));
   OS.write(static_cast<uint8_t>(Sym.SymInfo.Lang));
   writeVar(Strings.index(Sym.Name), OS);
   writeVar(Strings.index(Sym.Scope), OS);
std::pair<SymbolID, std::vector<Ref>> readRefs(Reader &Data,
   }
   return Result;
 }
 
 // FILE ENCODING
 // A file is a RIFF chunk with type 'CdIx'.
 // It contains the sections:
 //   - meta: version number
-//   - srcs: checksum of the source file
+//   - srcs: information related to source files
 //   - stri: string table
 //   - symb: symbols
 //   - refs: references to symbols
 
 // The current versioning scheme is simple - non-current versions are rejected.
 // If you make a breaking change, bump this version number to invalidate stored
 // data. Later we may want to support some backward compatibility.
 constexpr static uint32_t Version = 7;
if (Meta.consume32() != Version)
     return makeError("wrong version");
 
   auto Strings = readStringTable(Chunks.lookup("stri"));
   if (!Strings)
     return Strings.takeError();
 
   IndexFileIn Result;
   if (Chunks.count("srcs")) {
-    Reader Hash(Chunks.lookup("srcs"));
-    Result.Digest.emplace();
-    llvm::StringRef Digest = Hash.consume(Result.Digest->size());
-    std::copy(Digest.bytes_begin(), Digest.bytes_end(), Result.Digest->begin());
+    Reader SrcsReader(Chunks.lookup("srcs"));
+    Result.Sources.emplace();
+    while (!SrcsReader.eof()) {
+      auto SF = readSourceFile(SrcsReader, Strings->Strings);
+      auto Entry = Result.Sources->try_emplace(SF.FileURI).first;
+      Entry->getValue() = std::move(SF);
+      Entry->getValue().FileURI = Entry->getKey();
+      for(auto &Include : Entry->getValue().DirectIncludes)
+        Include = Result.Sources->try_emplace(Include).first->getKey();
+    }
+    if (SrcsReader.err())
+      return makeError("malformed or truncated include uri");
   }
 
   if (Chunks.count("symb")) {
     Reader SymbolReader(Chunks.lookup("symb"));
     SymbolSlab::Builder Symbols;
     while (!SymbolReader.eof())
       Symbols.insert(readSymbol(SymbolReader, Strings->Strings));
     if (SymbolReader.err())
if (Chunks.count("refs")) {
     }
     if (RefsReader.err())
       return makeError("malformed or truncated refs");
     Result.Refs = std::move(Refs).build();
   }
   return std::move(Result);
 }
 
+template <class Callback>
+void visitStrings(IndexFileIn::SourceFile &SF, const Callback &CB) {
+  CB(SF.FileURI);
+  for (llvm::StringRef &Include : SF.DirectIncludes)
+    CB(Include);
+}
+
 void writeRIFF(const IndexFileOut &Data, raw_ostream &OS) {
   assert(Data.Symbols && "An index file without symbols makes no sense!");
   riff::File RIFF;
   RIFF.Type = riff::fourCC("CdIx");
 
   SmallString<4> Meta;
   {
     raw_svector_ostream MetaOS(Meta);
     write32(Version, MetaOS);
   }
   RIFF.Chunks.push_back({riff::fourCC("meta"), Meta});
 
-  if (Data.Digest) {
-    llvm::StringRef Hash(reinterpret_cast<const char *>(Data.Digest->data()),
-                         Data.Digest->size());
-    RIFF.Chunks.push_back({riff::fourCC("srcs"), Hash});
-  }
-
   StringTableOut Strings;
   std::vector<Symbol> Symbols;
   for (const auto &Sym : *Data.Symbols) {
     Symbols.emplace_back(Sym);
     visitStrings(Symbols.back(), [&](StringRef &S) { Strings.intern(S); });
   }
+  std::vector<IndexFileIn::SourceFile> Sources;
+  if (Data.Sources)
+    for (const auto &Source : *Data.Sources) {
+      Sources.push_back(Source.getValue());
+      visitStrings(Sources.back(), [&](StringRef &S) { Strings.intern(S); });
+    }
+
   std::vector<std::pair<SymbolID, std::vector<Ref>>> Refs;
   if (Data.Refs) {
     for (const auto &Sym : *Data.Refs) {
       Refs.emplace_back(Sym);
       for (auto &Ref : Refs.back().second) {
         StringRef File = Ref.Location.FileURI;
         Strings.intern(File);
         Ref.Location.FileURI = File.data();
if (Data.Refs) {
     {
       raw_string_ostream RefsOS(RefsSection);
       for (const auto &Sym : Refs)
         writeRefs(Sym.first, Sym.second, Strings, RefsOS);
     }
     RIFF.Chunks.push_back({riff::fourCC("refs"), RefsSection});
   }
 
+  std::string SrcsSection;
+  {
+    {
+      raw_string_ostream SrcsOS(SrcsSection);
+      for (const auto &SF : Sources)
+        writeSourceFile(SF, Strings, SrcsOS);
+    }
+    RIFF.Chunks.push_back({riff::fourCC("srcs"), SrcsSection});
+  }
+
   OS << RIFF;
 }
 
 } // namespace
 
 // Defined in YAMLSerialization.cpp.
 void writeYAML(const IndexFileOut &, raw_ostream &);
 Expected<IndexFileIn> readYAML(StringRef);

unittests/clangd/BackgroundIndexTests.cpp

TEST(BackgroundIndexTest, ShardStorageWriteTest) {
   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")}));
-  EXPECT_EQ(*ShardSource->Digest, Digest);
+  EXPECT_EQ(ShardSource->Sources->lookup("unittest:///root/A.cc").Digest,
+            Digest);
+}
+
+TEST(BackgroundIndexTest, DirectIncludesTest) {
+  MockFSProvider FS;
+  FS.Files[testPath("root/B.h")] = "";
+  FS.Files[testPath("root/A.h")] = R"cpp(
+      #include "B.h"
+      void common();
+      void f_b();
+      class A_CC {};
+      )cpp";
+  std::string A_CC = "#include \"A.h\"\nvoid g() { (void)common; }";
+  FS.Files[testPath("root/A.cc")] = A_CC;
+
+  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")};
+  {
+    BackgroundIndex Idx(Context::empty(), "", FS,
+                        [&](llvm::StringRef) { return &MSS; });
+    Idx.enqueue(testPath("root"), Cmd);
+    Idx.blockUntilIdleForTest();
+  }
+
+  auto ShardSource = MSS.loadShard(testPath("root/A.cc"));
+  EXPECT_TRUE(ShardSource->Sources);
+  EXPECT_EQ(ShardSource->Sources->size(), 3U); // A.cc, A.h, B.h
+  EXPECT_THAT(
+      ShardSource->Sources->lookup("unittest:///root/A.cc").DirectIncludes,
+      UnorderedElementsAre("unittest:///root/A.h"));
+
+  auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
+  EXPECT_TRUE(ShardHeader->Sources);
+  EXPECT_EQ(ShardHeader->Sources->size(), 2U); // A.h B.h
+  EXPECT_THAT(
+      ShardHeader->Sources->lookup("unittest:///root/A.h").DirectIncludes,
+      UnorderedElementsAre("unittest:///root/B.h"));
 }
 
 } // namespace clangd
 } // namespace clang

unittests/clangd/SerializationTests.cpp

TEST(SerializationTest, BinaryConversions) {
 
   // Assert the YAML serializations match, for nice comparisons and diffs.
   EXPECT_THAT(YAMLFromSymbols(*In2->Symbols),
               UnorderedElementsAreArray(YAMLFromSymbols(*In->Symbols)));
   EXPECT_THAT(YAMLFromRefs(*In2->Refs),
               UnorderedElementsAreArray(YAMLFromRefs(*In->Refs)));
 }
 
-TEST(SerializationTest, HashTest) {
+TEST(SerializationTest, SrcsTest) {
   auto In = readIndexFile(YAML);
   EXPECT_TRUE(bool(In)) << In.takeError();
 
-  std::string TestContent("TESTCONTENT");
-  auto Digest =
+  std::string TestContent("TestContent");
+  IndexFileIn::SourceFile SF;
+  SF.Digest =
       llvm::SHA1::hash({reinterpret_cast<const uint8_t *>(TestContent.data()),
                         TestContent.size()});
+  SF.DirectIncludes = {"inc1", "inc2"};
+  SF.FileURI = "FileURI";
+  SF.IsTU = true;
+  llvm::StringMap<IndexFileIn::SourceFile> Sources;
+  Sources[SF.FileURI] = SF;
   // Write to binary format, and parse again.
   IndexFileOut Out(*In);
   Out.Format = IndexFileFormat::RIFF;
-  Out.Digest = &Digest;
+  Out.Sources = &Sources;
+  {
     std::string Serialized = to_string(Out);
 
-  auto In2 = readIndexFile(Serialized);
-  ASSERT_TRUE(bool(In2)) << In.takeError();
-  ASSERT_EQ(In2->Digest, Digest);
-  ASSERT_TRUE(In2->Symbols);
-  ASSERT_TRUE(In2->Refs);
-
+    auto In = readIndexFile(Serialized);
+    ASSERT_TRUE(bool(In)) << In.takeError();
+    ASSERT_TRUE(In->Symbols);
+    ASSERT_TRUE(In->Refs);
+    ASSERT_TRUE(In->Sources);
+    ASSERT_TRUE(In->Sources->count(SF.FileURI));
     // Assert the YAML serializations match, for nice comparisons and diffs.
-  EXPECT_THAT(YAMLFromSymbols(*In2->Symbols),
+    EXPECT_THAT(YAMLFromSymbols(*In->Symbols),
                 UnorderedElementsAreArray(YAMLFromSymbols(*In->Symbols)));
-  EXPECT_THAT(YAMLFromRefs(*In2->Refs),
+    EXPECT_THAT(YAMLFromRefs(*In->Refs),
                 UnorderedElementsAreArray(YAMLFromRefs(*In->Refs)));
+    auto SFDeserialized = In->Sources->lookup(SF.FileURI);
+    EXPECT_EQ(SFDeserialized.Digest, SF.Digest);
+    EXPECT_EQ(SFDeserialized.DirectIncludes, SF.DirectIncludes);
+    EXPECT_EQ(SFDeserialized.FileURI, SF.FileURI);
+    EXPECT_EQ(SFDeserialized.IsTU, SF.IsTU);
+  }
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang