1 files changed, 30 insertions, 2 deletions
diff --git a/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp b/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp
index bb0e730edd3..f76cb6ae5f6 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/DbiStreamBuilder.cpp
@@ -9,6 +9,7 @@
 
 #include "llvm/DebugInfo/PDB/Raw/DbiStreamBuilder.h"
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/DebugInfo/MSF/MSFBuilder.h"
 #include "llvm/DebugInfo/MSF/MappedBlockStream.h"
 #include "llvm/DebugInfo/MSF/StreamWriter.h"
@@ -43,10 +44,24 @@ void DbiStreamBuilder::setFlags(uint16_t F) { Flags = F; }
 
 void DbiStreamBuilder::setMachineType(PDB_Machine M) { MachineType = M; }
 
+Error DbiStreamBuilder::addDbgStream(pdb::DbgHeaderType Type,
+                                     ArrayRef<uint8_t> Data) {
+  if (DbgStreams[(int)Type].StreamNumber == kInvalidStreamIndex)
+    return make_error<RawError>(raw_error_code::duplicate_entry,
+                                "The specified stream type already exists");
+  auto ExpectedIndex = Msf.addStream(Data.size());
+  if (!ExpectedIndex)
+    return ExpectedIndex.takeError();
+  uint32_t Index = std::move(*ExpectedIndex);
+  DbgStreams[(int)Type].Data = Data;
+  DbgStreams[(int)Type].StreamNumber = Index;
+  return Error::success();
+}
+
 uint32_t DbiStreamBuilder::calculateSerializedLength() const {
   // For now we only support serializing the header.
   return sizeof(DbiStreamHeader) + calculateFileInfoSubstreamSize() +
-         calculateModiSubstreamSize();
+         calculateModiSubstreamSize() + DbgStreams.size() * sizeof(uint16_t);
 }
 
 Error DbiStreamBuilder::addModuleInfo(StringRef ObjFile, StringRef Module) {
@@ -216,7 +231,7 @@ Error DbiStreamBuilder::finalize() {
   H->ECSubstreamSize = 0;
   H->FileInfoSize = FileInfoBuffer.getLength();
   H->ModiSubstreamSize = ModInfoBuffer.getLength();
-  H->OptionalDbgHdrSize = 0;
+  H->OptionalDbgHdrSize = DbgStreams.size() * sizeof(uint16_t);
   H->SecContrSubstreamSize = 0;
   H->SectionMapSize = 0;
   H->TypeServerSize = 0;
@@ -273,6 +288,19 @@ Error DbiStreamBuilder::commit(const msf::MSFLayout &Layout,
     return EC;
   if (auto EC = Writer.writeStreamRef(FileInfoBuffer))
     return EC;
+  for (auto &Stream : DbgStreams)
+    if (auto EC = Writer.writeInteger(Stream.StreamNumber))
+      return EC;
+
+  for (auto &Stream : DbgStreams) {
+    if (Stream.StreamNumber == kInvalidStreamIndex)
+      continue;
+    auto WritableStream = WritableMappedBlockStream::createIndexedStream(
+        Layout, Buffer, Stream.StreamNumber);
+    StreamWriter DbgStreamWriter(*WritableStream);
+    if (auto EC = DbgStreamWriter.writeArray(Stream.Data))
+      return EC;
+  }
 
   if (Writer.bytesRemaining() > 0)
     return make_error<RawError>(raw_error_code::invalid_format,