[Serialization] Serialize/deserialize source file list

Author: rintaro

include/swift/AST/FileUnit.h

@@ -402,6 +402,9 @@ class LoadedFile : public FileUnit {
     return false;
   }
 
+  virtual void
+  collectSourceFileNames(llvm::function_ref<void(StringRef)> callback) const {}
+
   static bool classof(const FileUnit *file) {
     return file->getKind() == FileUnitKind::SerializedAST ||
            file->getKind() == FileUnitKind::ClangModule ||

include/swift/AST/Module.h

@@ -728,6 +728,8 @@ class ModuleDecl : public DeclContext, public TypeDecl {
     return ReverseFullNameIterator(this);
   }
 
+  void collectSourceFileNames(llvm::function_ref<void(StringRef)> callback);
+
   SourceRange getSourceRange() const { return SourceRange(); }
 
   static bool classof(const DeclContext *DC) {

include/swift/Serialization/SerializedModuleLoader.h

@@ -446,6 +446,9 @@ class SerializedASTFile final : public LoadedFile {
 
   StringRef getTargetTriple() const;
 
+  virtual void collectSourceFileNames(
+      llvm::function_ref<void(StringRef)>) const override;
+
   static bool classof(const FileUnit *file) {
     return file->getKind() == FileUnitKind::SerializedAST;
   }

lib/AST/Module.cpp

@@ -1514,6 +1514,17 @@ const clang::Module *ModuleDecl::findUnderlyingClangModule() const {
   return nullptr;
 }
 
+void ModuleDecl::collectSourceFileNames(
+    llvm::function_ref<void(StringRef)> callback) {
+  for (FileUnit *fileUnit : getFiles()) {
+    if (SourceFile *SF = dyn_cast<SourceFile>(fileUnit)) {
+      callback(SF->getFilename());
+    } else if (auto *serialized = dyn_cast<LoadedFile>(fileUnit)) {
+      serialized->collectSourceFileNames(callback);
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Cross-Import Overlays
 //===----------------------------------------------------------------------===//

lib/Serialization/ModuleFile.cpp

@@ -959,6 +959,24 @@ Optional<CommentInfo> ModuleFile::getCommentForDecl(const Decl *D) const {
   return getCommentForDeclByUSR(USRBuffer.str());
 }
 
+void ModuleFile::collectSourceFileNames(
+    llvm::function_ref<void(StringRef)> callback) const {
+  if (Core->SourceFileListData.empty())
+    return;
+  assert(!Core->SourceLocsTextData.empty());
+
+  auto *Cursor = Core->SourceFileListData.bytes_begin();
+  auto *End = Core->SourceFileListData.bytes_end();
+  while (Cursor < End) {
+    auto fileID = endian::readNext<uint32_t, little, unaligned>(Cursor);
+    assert(fileID < Core->SourceLocsTextData.size());
+    auto filePath = Core->SourceLocsTextData.substr(fileID);
+    size_t terminatorOffset = filePath.find('\0');
+    filePath = filePath.slice(0, terminatorOffset);
+    callback(filePath);
+  }
+}
+
 Optional<BasicDeclLocs>
 ModuleFile::getBasicDeclLocsForDecl(const Decl *D) const {
   assert(D);

lib/Serialization/ModuleFile.h

@@ -694,6 +694,7 @@ class ModuleFile
   Optional<BasicDeclLocs> getBasicDeclLocsForDecl(const Decl *D) const;
   Identifier getDiscriminatorForPrivateValue(const ValueDecl *D);
   Optional<Fingerprint> loadFingerprint(const IterableDeclContext *IDC) const;
+  void collectSourceFileNames(llvm::function_ref<void(StringRef)> callback) const;
 
 
   // MARK: Deserialization interface

lib/Serialization/ModuleFileSharedCore.cpp

@@ -974,6 +974,9 @@ bool ModuleFileSharedCore::readDeclLocsBlock(llvm::BitstreamCursor &cursor) {
         return false;
       }
       switch (*kind) {
+      case decl_locs_block::SOURCE_FILE_LIST:
+        SourceFileListData = blobData;
+        break;
       case decl_locs_block::BASIC_DECL_LOCS:
         BasicDeclLocsData = blobData;
         break;

lib/Serialization/ModuleFileSharedCore.h

@@ -289,6 +289,9 @@ class ModuleFileSharedCore {
   /// A blob of 0 terminated string segments referenced in \c SourceLocsTextData
   StringRef SourceLocsTextData;
 
+  /// A blob of source file list.
+  StringRef SourceFileListData;
+
   /// An array of fixed size source location data for each USR appearing in
   /// \c DeclUSRsTable.
   StringRef BasicDeclLocsData;

lib/Serialization/SerializeDoc.cpp

@@ -783,6 +783,46 @@ static void emitBasicLocsRecord(llvm::BitstreamWriter &Out,
   DeclLocsList.emit(scratch, Writer.Buffer);
 }
 
+static void emitFileListRecord(llvm::BitstreamWriter &Out,
+                               ModuleOrSourceFile MSF, StringWriter &FWriter) {
+  assert(MSF);
+
+  struct SourceFileListWriter {
+    StringWriter &FWriter;
+
+    llvm::SmallString<1024> Buffer;
+    llvm::StringSet<> seenFilenames;
+
+    void emitFilename(StringRef filename) {
+      if (filename.empty())
+        return;
+      llvm::SmallString<128> absolutePath = filename;
+      llvm::sys::fs::make_absolute(absolutePath);
+      if (!seenFilenames.insert(absolutePath).second)
+        return;
+
+      auto fileID = FWriter.getTextOffset(absolutePath);
+      llvm::raw_svector_ostream out(Buffer);
+      endian::Writer writer(out, little);
+      writer.write<uint32_t>(fileID);
+    }
+
+    SourceFileListWriter(StringWriter &FWriter) : FWriter(FWriter) {}
+  } writer(FWriter);
+
+  if (SourceFile *SF = MSF.dyn_cast<SourceFile *>()) {
+    writer.emitFilename(SF->getFilename());
+  } else {
+    auto *M = MSF.get<ModuleDecl *>();
+    M->collectSourceFileNames(
+        [&](StringRef filename) { writer.emitFilename(filename); });
+  }
+
+  const decl_locs_block::SourceFileListLayout FileList(Out);
+  SmallVector<uint64_t, 8> scratch;
+  FileList.emit(scratch, writer.Buffer);
+}
+
 class SourceInfoSerializer : public SerializerBase {
 public:
   using SerializerBase::SerializerBase;
@@ -807,6 +847,7 @@ class SourceInfoSerializer : public SerializerBase {
     BLOCK_RECORD(control_block, TARGET);
 
     BLOCK(DECL_LOCS_BLOCK);
+    BLOCK_RECORD(decl_locs_block, SOURCE_FILE_LIST);
     BLOCK_RECORD(decl_locs_block, BASIC_DECL_LOCS);
     BLOCK_RECORD(decl_locs_block, DECL_USRS);
     BLOCK_RECORD(decl_locs_block, TEXT_DATA);
@@ -849,6 +890,7 @@ void serialization::writeSourceInfoToStream(raw_ostream &os,
       DeclUSRsTableWriter USRWriter;
       StringWriter FPWriter;
       DocRangeWriter DocWriter;
+      emitFileListRecord(S.Out, DC, FPWriter);
       emitBasicLocsRecord(S.Out, DC, USRWriter, FPWriter, DocWriter);
       // Emit USR table mapping from a USR to USR Id.
       // The basic locs record uses USR Id instead of actual USR, so that we

lib/Serialization/SerializedModuleLoader.cpp

@@ -1327,3 +1327,8 @@ SerializedASTFile::getDiscriminatorForPrivateValue(const ValueDecl *D) const {
   assert(!discriminator.empty() && "no discriminator found for value");
   return discriminator;
 }
+
+void SerializedASTFile::collectSourceFileNames(
+    llvm::function_ref<void(StringRef)> callback) const {
+  File.collectSourceFileNames(callback);
+}