5 files changed, 226 insertions, 64 deletions

diff --git a/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp b/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp
index 3c0586c..63a7ba1 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/DbiStream.cpp
@@ -142,14 +142,11 @@ Error DbiStream::reload() {
     return make_error<RawError>(raw_error_code::corrupt_file,
                                 "DBI type server substream not aligned.");
 
-  // Since each ModInfo in the stream is a variable length, we have to iterate
-  // them to know how many there actually are.
-  VarStreamArray<ModInfo> ModInfoArray;
-  if (auto EC = Reader.readArray(ModInfoArray, Header->ModiSubstreamSize))
+  if (auto EC =
+          Reader.readStreamRef(ModInfoSubstream, Header->ModiSubstreamSize))
+    return EC;
+  if (auto EC = initializeModInfoArray())
     return EC;
-  for (auto &Info : ModInfoArray) {
-    ModuleInfos.emplace_back(Info);
-  }
 
   if (auto EC = Reader.readStreamRef(SecContrSubstream,
                                      Header->SecContrSubstreamSize))
@@ -285,6 +282,24 @@ Error DbiStream::initializeSectionContributionData() {
                               "Unsupported DBI Section Contribution version");
 }
 
+Error DbiStream::initializeModInfoArray() {
+  if (ModInfoSubstream.getLength() == 0)
+    return Error::success();
+
+  // Since each ModInfo in the stream is a variable length, we have to iterate
+  // them to know how many there actually are.
+  StreamReader Reader(ModInfoSubstream);
+
+  VarStreamArray<ModInfo> ModInfoArray;
+  if (auto EC = Reader.readArray(ModInfoArray, ModInfoSubstream.getLength()))
+    return EC;
+  for (auto &Info : ModInfoArray) {
+    ModuleInfos.emplace_back(Info);
+  }
+
+  return Error::success();
+}
+
 // Initializes this->SectionHeaders.
 Error DbiStream::initializeSectionHeadersData() {
   if (DbgStreams.size() == 0)
@@ -437,7 +452,10 @@ Error DbiStream::initializeFileInfo() {
 }
 
 uint32_t DbiStream::getDebugStreamIndex(DbgHeaderType Type) const {
-  return DbgStreams[static_cast<uint16_t>(Type)];
+  uint16_t T = static_cast<uint16_t>(Type);
+  if (T >= DbgStreams.size())
+    return DbiStream::InvalidStreamIndex;
+  return DbgStreams[T];
 }
 
 Expected<StringRef> DbiStream::getFileNameForIndex(uint32_t Index) const {
@@ -458,5 +476,26 @@ Error DbiStream::commit() {
   if (auto EC = Writer.writeObject(*Header))
     return EC;
 
+  if (auto EC = Writer.writeStreamRef(ModInfoSubstream))
+    return EC;
+
+  if (auto EC = Writer.writeStreamRef(SecContrSubstream,
+                                      SecContrSubstream.getLength()))
+    return EC;
+  if (auto EC =
+          Writer.writeStreamRef(SecMapSubstream, SecMapSubstream.getLength()))
+    return EC;
+  if (auto EC = Writer.writeStreamRef(FileInfoSubstream,
+                                      FileInfoSubstream.getLength()))
+    return EC;
+  if (auto EC = Writer.writeStreamRef(TypeServerMapSubstream,
+                                      TypeServerMapSubstream.getLength()))
+    return EC;
+  if (auto EC = Writer.writeStreamRef(ECSubstream, ECSubstream.getLength()))
+    return EC;
+
+  if (Writer.bytesRemaining() > 0)
+    return make_error<RawError>(raw_error_code::invalid_format,
+                                "Unexpected bytes found in DBI Stream");
   return Error::success();
 }
diff --git a/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp b/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp
index 34ff8ae..bbb0876 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp
@@ -18,9 +18,13 @@ using namespace llvm;
 using namespace llvm::codeview;
 using namespace llvm::pdb;
 
-DbiStreamBuilder::DbiStreamBuilder()
-    : Age(1), BuildNumber(0), PdbDllVersion(0), PdbDllRbld(0), Flags(0),
-      MachineType(PDB_Machine::x86) {}
+namespace {
+class ModiSubstreamBuilder {};
+}
+
+DbiStreamBuilder::DbiStreamBuilder(BumpPtrAllocator &Allocator)
+    : Allocator(Allocator), Age(1), BuildNumber(0), PdbDllVersion(0),
+      PdbDllRbld(0), Flags(0), MachineType(PDB_Machine::x86) {}
 
 void DbiStreamBuilder::setVersionHeader(PdbRaw_DbiVer V) { VerHeader = V; }
 
@@ -38,7 +42,157 @@ void DbiStreamBuilder::setMachineType(PDB_Machine M) { MachineType = M; }
 
 uint32_t DbiStreamBuilder::calculateSerializedLength() const {
   // For now we only support serializing the header.
-  return sizeof(DbiStream::HeaderInfo);
+  return sizeof(DbiStream::HeaderInfo) + calculateFileInfoSubstreamSize() +
+         calculateModiSubstreamSize();
+}
+
+Error DbiStreamBuilder::addModuleInfo(StringRef ObjFile, StringRef Module) {
+  auto Entry = llvm::make_unique<ModuleInfo>();
+  ModuleInfo *M = Entry.get();
+  Entry->Mod = Module;
+  Entry->Obj = ObjFile;
+  auto Result = ModuleInfos.insert(std::make_pair(Module, std::move(Entry)));
+  if (!Result.second)
+    return make_error<RawError>(raw_error_code::duplicate_entry,
+                                "The specified module already exists");
+  ModuleInfoList.push_back(M);
+  return Error::success();
+}
+
+Error DbiStreamBuilder::addModuleSourceFile(StringRef Module, StringRef File) {
+  auto ModIter = ModuleInfos.find(Module);
+  if (ModIter == ModuleInfos.end())
+    return make_error<RawError>(raw_error_code::no_entry,
+                                "The specified module was not found");
+  uint32_t Index = SourceFileNames.size();
+  SourceFileNames.insert(std::make_pair(File, Index));
+  auto &ModEntry = *ModIter;
+  ModEntry.second->SourceFiles.push_back(File);
+  return Error::success();
+}
+
+uint32_t DbiStreamBuilder::calculateModiSubstreamSize() const {
+  uint32_t Size = 0;
+  for (const auto &M : ModuleInfoList) {
+    Size += sizeof(ModInfo::FileLayout);
+    Size += M->Mod.size() + 1;
+    Size += M->Obj.size() + 1;
+  }
+  return Size;
+}
+
+uint32_t DbiStreamBuilder::calculateFileInfoSubstreamSize() const {
+  //   ulittle16_t NumModules;
+  //   ulittle16_t NumSourceFiles;
+  //   ulittle16_t ModIndices[NumModules];
+  //   ulittle16_t ModFileCounts[NumModules];
+  //   ulittle32_t FileNameOffsets[NumSourceFiles];
+  //   char Names[NumSourceFiles][];
+  uint32_t Size = 0;
+  Size += sizeof(ulittle16_t);                         // NumModules
+  Size += sizeof(ulittle16_t);                         // NumSourceFiles
+  Size += ModuleInfoList.size() * sizeof(ulittle16_t); // ModIndices
+  Size += ModuleInfoList.size() * sizeof(ulittle16_t); // ModFileCounts
+  uint32_t NumFileInfos = 0;
+  for (const auto &M : ModuleInfoList)
+    NumFileInfos += M->SourceFiles.size();
+  Size += NumFileInfos * sizeof(ulittle32_t); // FileNameOffsets
+  Size += calculateNamesBufferSize();
+  return Size;
+}
+
+uint32_t DbiStreamBuilder::calculateNamesBufferSize() const {
+  uint32_t Size = 0;
+  for (const auto &F : SourceFileNames) {
+    Size += F.getKeyLength() + 1; // Names[I];
+  }
+  return Size;
+}
+
+Error DbiStreamBuilder::generateModiSubstream() {
+  uint32_t Size = calculateModiSubstreamSize();
+  auto Data = Allocator.Allocate<uint8_t>(Size);
+
+  ModInfoBuffer = ByteStream<true>(MutableArrayRef<uint8_t>(Data, Size));
+
+  StreamWriter ModiWriter(ModInfoBuffer);
+  for (const auto &M : ModuleInfoList) {
+    ModInfo::FileLayout Layout = {};
+    Layout.ModDiStream = DbiStream::InvalidStreamIndex;
+    Layout.NumFiles = M->SourceFiles.size();
+    if (auto EC = ModiWriter.writeObject(Layout))
+      return EC;
+    if (auto EC = ModiWriter.writeZeroString(M->Mod))
+      return EC;
+    if (auto EC = ModiWriter.writeZeroString(M->Obj))
+      return EC;
+  }
+  if (ModiWriter.bytesRemaining() != 0)
+    return make_error<RawError>(raw_error_code::invalid_format,
+                                "Unexpected bytes in Modi Stream Data");
+  return Error::success();
+}
+
+Error DbiStreamBuilder::generateFileInfoSubstream() {
+  uint32_t Size = calculateFileInfoSubstreamSize();
+  uint32_t NameSize = calculateNamesBufferSize();
+  auto Data = Allocator.Allocate<uint8_t>(Size);
+  uint32_t NamesOffset = Size - NameSize;
+
+  FileInfoBuffer = ByteStream<true>(MutableArrayRef<uint8_t>(Data, Size));
+
+  StreamRef MetadataBuffer = StreamRef(FileInfoBuffer).keep_front(NamesOffset);
+  StreamWriter MetadataWriter(MetadataBuffer);
+
+  uint16_t ModiCount = std::min<uint16_t>(UINT16_MAX, ModuleInfos.size());
+  uint16_t FileCount = std::min<uint16_t>(UINT16_MAX, SourceFileNames.size());
+  if (auto EC = MetadataWriter.writeInteger(ModiCount)) // NumModules
+    return EC;
+  if (auto EC = MetadataWriter.writeInteger(FileCount)) // NumSourceFiles
+    return EC;
+  for (uint16_t I = 0; I < ModiCount; ++I) {
+    if (auto EC = MetadataWriter.writeInteger(I)) // Mod Indices
+      return EC;
+  }
+  for (const auto MI : ModuleInfoList) {
+    FileCount = static_cast<uint16_t>(MI->SourceFiles.size());
+    if (auto EC = MetadataWriter.writeInteger(FileCount)) // Mod File Counts
+      return EC;
+  }
+
+  // Before writing the FileNameOffsets array, write the NamesBuffer array.
+  // A side effect of this is that this will actually compute the various
+  // file name offsets, so we can then go back and write the FileNameOffsets
+  // array to the other substream.
+  NamesBuffer = StreamRef(FileInfoBuffer).drop_front(NamesOffset);
+  StreamWriter NameBufferWriter(NamesBuffer);
+  for (auto &Name : SourceFileNames) {
+    Name.second = NameBufferWriter.getOffset();
+    if (auto EC = NameBufferWriter.writeZeroString(Name.getKey()))
+      return EC;
+  }
+
+  for (const auto MI : ModuleInfoList) {
+    for (StringRef Name : MI->SourceFiles) {
+      auto Result = SourceFileNames.find(Name);
+      if (Result == SourceFileNames.end())
+        return make_error<RawError>(raw_error_code::no_entry,
+                                    "The source file was not found.");
+      if (auto EC = MetadataWriter.writeInteger(Result->second))
+        return EC;
+    }
+  }
+
+  if (NameBufferWriter.bytesRemaining() > 0)
+    return make_error<RawError>(raw_error_code::invalid_format,
+                                "The names buffer contained unexpected data.");
+
+  if (MetadataWriter.bytesRemaining() > 0)
+    return make_error<RawError>(
+        raw_error_code::invalid_format,
+        "The metadata buffer contained unexpected data.");
+
+  return Error::success();
 }
 
 Expected<std::unique_ptr<DbiStream>> DbiStreamBuilder::build(PDBFile &File) {
@@ -54,6 +208,12 @@ Expected<std::unique_ptr<DbiStream>> DbiStreamBuilder::build(PDBFile &File) {
       static_cast<DbiStream::HeaderInfo *>(DS->getAllocator().Allocate(
           sizeof(DbiStream::HeaderInfo),
           llvm::AlignOf<DbiStream::HeaderInfo>::Alignment));
+
+  if (auto EC = generateModiSubstream())
+    return std::move(EC);
+  if (auto EC = generateFileInfoSubstream())
+    return std::move(EC);
+
   H->VersionHeader = *VerHeader;
   H->VersionSignature = -1;
   H->Age = Age;
@@ -64,8 +224,8 @@ Expected<std::unique_ptr<DbiStream>> DbiStreamBuilder::build(PDBFile &File) {
   H->MachineType = static_cast<uint16_t>(MachineType);
 
   H->ECSubstreamSize = 0;
-  H->FileInfoSize = 0;
-  H->ModiSubstreamSize = 0;
+  H->FileInfoSize = FileInfoBuffer.getLength();
+  H->ModiSubstreamSize = ModInfoBuffer.getLength();
   H->OptionalDbgHdrSize = 0;
   H->SecContrSubstreamSize = 0;
   H->SectionMapSize = 0;
@@ -77,5 +237,11 @@ Expected<std::unique_ptr<DbiStream>> DbiStreamBuilder::build(PDBFile &File) {
 
   auto Dbi = llvm::make_unique<DbiStream>(File, std::move(DS));
   Dbi->Header = H;
+  Dbi->FileInfoSubstream = StreamRef(FileInfoBuffer);
+  Dbi->ModInfoSubstream = StreamRef(ModInfoBuffer);
+  if (auto EC = Dbi->initializeModInfoArray())
+    return std::move(EC);
+  if (auto EC = Dbi->initializeFileInfo())
+    return std::move(EC);
   return std::move(Dbi);
 }
diff --git a/llvm/lib/DebugInfo/PDB/Raw/ModInfo.cpp b/llvm/lib/DebugInfo/PDB/Raw/ModInfo.cpp
index bae135f..538fe22 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/ModInfo.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/ModInfo.cpp
@@ -17,55 +17,6 @@ using namespace llvm;
 using namespace llvm::pdb;
 using namespace llvm::support;
 
-namespace {
-
-struct SCBytes {
-  ulittle16_t Section;
-  char Padding1[2];
-  little32_t Offset;
-  little32_t Size;
-  ulittle32_t Characteristics;
-  ulittle16_t ModuleIndex;
-  char Padding2[2];
-  ulittle32_t DataCrc;
-  ulittle32_t RelocCrc;
-};
-
-// struct Flags {
-//  uint16_t fWritten : 1;   // True if ModInfo is dirty
-//  uint16_t fECEnabled : 1; // Is EC symbolic info present?  (What is EC?)
-//  uint16_t unused : 6;     // Reserved
-//  uint16_t iTSM : 8;       // Type Server Index for this module
-//};
-const uint16_t HasECFlagMask = 0x2;
-
-const uint16_t TypeServerIndexMask = 0xFF00;
-const uint16_t TypeServerIndexShift = 8;
-}
-
-struct ModInfo::FileLayout {
-  ulittle32_t Mod;           // Currently opened module.  This field is a
-                             // pointer in the reference implementation, but
-                             // that won't work on 64-bit systems, and anyway
-                             // it doesn't make sense to read a pointer from a
-                             // file.  For now it is unused, so just ignore it.
-  SCBytes SC;                // First section contribution of this module.
-  ulittle16_t Flags;         // See Flags definition.
-  ulittle16_t ModDiStream;   // Stream Number of module debug info
-  ulittle32_t SymBytes;      // Size of local symbol debug info in above stream
-  ulittle32_t LineBytes;     // Size of line number debug info in above stream
-  ulittle32_t C13Bytes;      // Size of C13 line number info in above stream
-  ulittle16_t NumFiles;      // Number of files contributing to this module
-  char Padding1[2];          // Padding so the next field is 4-byte aligned.
-  ulittle32_t FileNameOffs;  // array of [0..NumFiles) DBI name buffer offsets.
-                             // This field is a pointer in the reference
-                             // implementation, but as with `Mod`, we ignore it
-                             // for now since it is unused.
-  ulittle32_t SrcFileNameNI; // Name Index for src file name
-  ulittle32_t PdbFilePathNI; // Name Index for path to compiler PDB
-                             // Null terminated Module name
-                             // Null terminated Obj File Name
-};
 
 ModInfo::ModInfo() : Layout(nullptr) {}
 
diff --git a/llvm/lib/DebugInfo/PDB/Raw/PDBFileBuilder.cpp b/llvm/lib/DebugInfo/PDB/Raw/PDBFileBuilder.cpp
index 9063fd6..7ed7abef 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/PDBFileBuilder.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/PDBFileBuilder.cpp
@@ -53,7 +53,7 @@ InfoStreamBuilder &PDBFileBuilder::getInfoBuilder() {
 
 DbiStreamBuilder &PDBFileBuilder::getDbiBuilder() {
   if (!Dbi)
-    Dbi = llvm::make_unique<DbiStreamBuilder>();
+    Dbi = llvm::make_unique<DbiStreamBuilder>(File->Allocator);
   return *Dbi;
 }
 
diff --git a/llvm/lib/DebugInfo/PDB/Raw/RawError.cpp b/llvm/lib/DebugInfo/PDB/Raw/RawError.cpp
index eb169f7..cbce890 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/RawError.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/RawError.cpp
@@ -19,6 +19,8 @@ public:
       return "An unknown error has occurred.";
     case raw_error_code::feature_unsupported:
       return "The feature is unsupported by the implementation.";
+    case raw_error_code::invalid_format:
+      return "The record is in an unexpected format.";
     case raw_error_code::corrupt_file:
       return "The PDB file is corrupt.";
     case raw_error_code::insufficient_buffer:
@@ -30,6 +32,10 @@ public:
       return "The specified item does not exist in the array.";
     case raw_error_code::invalid_block_address:
       return "The specified block address is not valid.";
+    case raw_error_code::duplicate_entry:
+      return "The entry already exists.";
+    case raw_error_code::no_entry:
+      return "The entry does not exist.";
     case raw_error_code::not_writable:
       return "The PDB does not support writing.";
     case raw_error_code::invalid_tpi_hash: