1 files changed, 64 insertions, 11 deletions

diff --git a/clang-tools-extra/clangd/index/Serialization.cpp b/clang-tools-extra/clangd/index/Serialization.cpp
index 9eb4d21eb8f..1d17a42c486 100644
--- a/clang-tools-extra/clangd/index/Serialization.cpp
+++ b/clang-tools-extra/clangd/index/Serialization.cpp
@@ -247,6 +247,31 @@ SymbolLocation readLocation(Reader &Data, ArrayRef<StringRef> Strings) {
   return Loc;
 }
 
+IncludeGraphNode readIncludeGraphNode(Reader &Data,
+                                      llvm::ArrayRef<llvm::StringRef> Strings) {
+  IncludeGraphNode IGN;
+  IGN.IsTU = Data.consume8();
+  IGN.URI = Data.consumeString(Strings);
+  llvm::StringRef Digest = Data.consume(IGN.Digest.size());
+  std::copy(Digest.bytes_begin(), Digest.bytes_end(), IGN.Digest.begin());
+  IGN.DirectIncludes.resize(Data.consumeVar());
+  for (llvm::StringRef &Include : IGN.DirectIncludes)
+    Include = Data.consumeString(Strings);
+  return IGN;
+}
+
+void writeIncludeGraphNode(const IncludeGraphNode &IGN,
+                           const StringTableOut &Strings, raw_ostream &OS) {
+  OS.write(IGN.IsTU);
+  writeVar(Strings.index(IGN.URI), OS);
+  llvm::StringRef Hash(reinterpret_cast<const char *>(IGN.Digest.data()),
+                       IGN.Digest.size());
+  OS << Hash;
+  writeVar(IGN.DirectIncludes.size(), OS);
+  for (llvm::StringRef Include : IGN.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,
@@ -333,7 +358,7 @@ std::pair<SymbolID, std::vector<Ref>> readRefs(Reader &Data,
 // 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 include graph
 //   - stri: string table
 //   - symb: symbols
 //   - refs: references to symbols
@@ -367,10 +392,20 @@ Expected<IndexFileIn> readRIFF(StringRef Data) {
 
   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 IGN = readIncludeGraphNode(SrcsReader, Strings->Strings);
+      auto Entry = Result.Sources->try_emplace(IGN.URI).first;
+      Entry->getValue() = std::move(IGN);
+      // We change all the strings inside the structure to point at the keys in
+      // the map, since it is the only copy of the string that's going to live.
+      Entry->getValue().URI = 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")) {
@@ -397,6 +432,13 @@ Expected<IndexFileIn> readRIFF(StringRef Data) {
   return std::move(Result);
 }
 
+template <class Callback>
+void visitStrings(IncludeGraphNode &IGN, const Callback &CB) {
+  CB(IGN.URI);
+  for (llvm::StringRef &Include : IGN.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;
@@ -409,18 +451,19 @@ void writeRIFF(const IndexFileOut &Data, raw_ostream &OS) {
   }
   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<IncludeGraphNode> 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) {
@@ -458,6 +501,16 @@ void writeRIFF(const IndexFileOut &Data, raw_ostream &OS) {
     RIFF.Chunks.push_back({riff::fourCC("refs"), RefsSection});
   }
 
+  std::string SrcsSection;
+  {
+    {
+      raw_string_ostream SrcsOS(SrcsSection);
+      for (const auto &SF : Sources)
+        writeIncludeGraphNode(SF, Strings, SrcsOS);
+    }
+    RIFF.Chunks.push_back({riff::fourCC("srcs"), SrcsSection});
+  }
+
   OS << RIFF;
 }