1 files changed, 37 insertions, 16 deletions

diff --git a/lib/DebugInfo/PDB/Native/DbiStreamBuilder.cpp b/lib/DebugInfo/PDB/Native/DbiStreamBuilder.cpp
index c96553ff9b16..f6043bfd7cf9 100644
--- a/lib/DebugInfo/PDB/Native/DbiStreamBuilder.cpp
+++ b/lib/DebugInfo/PDB/Native/DbiStreamBuilder.cpp
@@ -27,7 +27,7 @@ using namespace llvm::pdb;
 DbiStreamBuilder::DbiStreamBuilder(msf::MSFBuilder &Msf)
     : Msf(Msf), Allocator(Msf.getAllocator()), Age(1), BuildNumber(0),
       PdbDllVersion(0), PdbDllRbld(0), Flags(0), MachineType(PDB_Machine::x86),
-      Header(nullptr), DbgStreams((int)DbgHeaderType::Max) {}
+      Header(nullptr) {}
 
 DbiStreamBuilder::~DbiStreamBuilder() {}
 
@@ -37,6 +37,14 @@ void DbiStreamBuilder::setAge(uint32_t A) { Age = A; }
 
 void DbiStreamBuilder::setBuildNumber(uint16_t B) { BuildNumber = B; }
 
+void DbiStreamBuilder::setBuildNumber(uint8_t Major, uint8_t Minor) {
+  BuildNumber = (uint16_t(Major) << DbiBuildNo::BuildMajorShift) &
+                DbiBuildNo::BuildMajorMask;
+  BuildNumber |= (uint16_t(Minor) << DbiBuildNo::BuildMinorShift) &
+                 DbiBuildNo::BuildMinorMask;
+  BuildNumber |= DbiBuildNo::NewVersionFormatMask;
+}
+
 void DbiStreamBuilder::setPdbDllVersion(uint16_t V) { PdbDllVersion = V; }
 
 void DbiStreamBuilder::setPdbDllRbld(uint16_t R) { PdbDllRbld = R; }
@@ -45,6 +53,11 @@ void DbiStreamBuilder::setFlags(uint16_t F) { Flags = F; }
 
 void DbiStreamBuilder::setMachineType(PDB_Machine M) { MachineType = M; }
 
+void DbiStreamBuilder::setMachineType(COFF::MachineTypes M) {
+  // These enums are mirrors of each other, so we can just cast the value.
+  MachineType = static_cast<pdb::PDB_Machine>(static_cast<unsigned>(M));
+}
+
 void DbiStreamBuilder::setSectionMap(ArrayRef<SecMapEntry> SecMap) {
   SectionMap = SecMap;
 }
@@ -63,15 +76,8 @@ void DbiStreamBuilder::setPublicsStreamIndex(uint32_t Index) {
 
 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;
+  DbgStreams[(int)Type].emplace();
+  DbgStreams[(int)Type]->Data = Data;
   return Error::success();
 }
 
@@ -258,7 +264,7 @@ Error DbiStreamBuilder::finalize() {
   H->TypeServerSize = 0;
   H->SymRecordStreamIndex = SymRecordStreamIndex;
   H->PublicSymbolStreamIndex = PublicsStreamIndex;
-  H->MFCTypeServerIndex = kInvalidStreamIndex;
+  H->MFCTypeServerIndex = 0; // Not sure what this is, but link.exe writes 0.
   H->GlobalSymbolStreamIndex = GlobalsStreamIndex;
 
   Header = H;
@@ -266,6 +272,15 @@ Error DbiStreamBuilder::finalize() {
 }
 
 Error DbiStreamBuilder::finalizeMsfLayout() {
+  for (auto &S : DbgStreams) {
+    if (!S.hasValue())
+      continue;
+    auto ExpectedIndex = Msf.addStream(S->Data.size());
+    if (!ExpectedIndex)
+      return ExpectedIndex.takeError();
+    S->StreamNumber = *ExpectedIndex;
+  }
+
   for (auto &MI : ModiList) {
     if (auto EC = MI->finalizeMsfLayout())
       return EC;
@@ -375,17 +390,23 @@ Error DbiStreamBuilder::commit(const msf::MSFLayout &Layout,
   if (auto EC = ECNamesBuilder.commit(Writer))
     return EC;
 
-  for (auto &Stream : DbgStreams)
-    if (auto EC = Writer.writeInteger(Stream.StreamNumber))
+  for (auto &Stream : DbgStreams) {
+    uint16_t StreamNumber = kInvalidStreamIndex;
+    if (Stream.hasValue())
+      StreamNumber = Stream->StreamNumber;
+    if (auto EC = Writer.writeInteger(StreamNumber))
       return EC;
+  }
 
   for (auto &Stream : DbgStreams) {
-    if (Stream.StreamNumber == kInvalidStreamIndex)
+    if (!Stream.hasValue())
       continue;
+    assert(Stream->StreamNumber != kInvalidStreamIndex);
+
     auto WritableStream = WritableMappedBlockStream::createIndexedStream(
-        Layout, MsfBuffer, Stream.StreamNumber, Allocator);
+        Layout, MsfBuffer, Stream->StreamNumber, Allocator);
     BinaryStreamWriter DbgStreamWriter(*WritableStream);
-    if (auto EC = DbgStreamWriter.writeArray(Stream.Data))
+    if (auto EC = DbgStreamWriter.writeArray(Stream->Data))
       return EC;
   }