1 files changed, 91 insertions, 3 deletions

diff --git a/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp b/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp
index cf0ef688b66..2d8e0d0ed91 100644
--- a/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp
+++ b/llvm/lib/DebugInfo/PDB/Raw/PDBFile.cpp
@@ -89,7 +89,7 @@ Error PDBFile::setBlockData(uint32_t BlockIndex, uint32_t Offset,
     return make_error<RawError>(
         raw_error_code::invalid_block_address,
         "setBlockData attempted to write out of block bounds.");
-  if (Data.size() >= getBlockSize() - Offset)
+  if (Data.size() > getBlockSize() - Offset)
     return make_error<RawError>(
         raw_error_code::invalid_block_address,
         "setBlockData attempted to write out of block bounds.");
@@ -167,6 +167,17 @@ llvm::ArrayRef<support::ulittle32_t> PDBFile::getDirectoryBlockArray() const {
   return DirectoryBlocks;
 }
 
+Expected<InfoStream &> PDBFile::emplacePDBInfoStream() {
+  if (Info)
+    Info.reset();
+
+  auto InfoS = MappedBlockStream::createIndexedStream(StreamPDB, *this);
+  if (!InfoS)
+    return InfoS.takeError();
+  Info = llvm::make_unique<InfoStream>(std::move(*InfoS));
+  return *Info;
+}
+
 Expected<InfoStream &> PDBFile::getPDBInfoStream() {
   if (!Info) {
     auto InfoS = MappedBlockStream::createIndexedStream(StreamPDB, *this);
@@ -340,12 +351,59 @@ void PDBFile::setStreamSizes(ArrayRef<support::ulittle32_t> Sizes) {
 }
 
 void PDBFile::setStreamMap(
-    ArrayRef<support::ulittle32_t> Directory,
     std::vector<ArrayRef<support::ulittle32_t>> &Streams) {
-  DirectoryBlocks = Directory;
   StreamMap = Streams;
 }
 
+void PDBFile::setDirectoryBlocks(ArrayRef<support::ulittle32_t> Directory) {
+  DirectoryBlocks = Directory;
+}
+
+Error PDBFile::generateSimpleStreamMap() {
+  if (StreamSizes.empty())
+    return Error::success();
+
+  static std::vector<std::vector<support::ulittle32_t>> StaticMap;
+  StreamMap.clear();
+  StaticMap.clear();
+
+  // Figure out how many blocks are needed for all streams, and set the first
+  // used block to the highest block so that we can write the rest of the
+  // blocks contiguously.
+  uint32_t TotalFileBlocks = getBlockCount();
+  std::vector<support::ulittle32_t> ReservedBlocks;
+  ReservedBlocks.push_back(support::ulittle32_t(0));
+  ReservedBlocks.push_back(SB->BlockMapAddr);
+  ReservedBlocks.insert(ReservedBlocks.end(), DirectoryBlocks.begin(),
+                        DirectoryBlocks.end());
+
+  uint32_t BlocksNeeded = 0;
+  for (auto Size : StreamSizes)
+    BlocksNeeded += bytesToBlocks(Size, getBlockSize());
+
+  support::ulittle32_t NextBlock(TotalFileBlocks - BlocksNeeded -
+                                 ReservedBlocks.size());
+
+  StaticMap.resize(StreamSizes.size());
+  for (uint32_t S = 0; S < StreamSizes.size(); ++S) {
+    uint32_t Size = StreamSizes[S];
+    uint32_t NumBlocks = bytesToBlocks(Size, getBlockSize());
+    auto &ThisStream = StaticMap[S];
+    for (uint32_t I = 0; I < NumBlocks;) {
+      NextBlock += 1;
+      if (std::find(ReservedBlocks.begin(), ReservedBlocks.end(), NextBlock) !=
+          ReservedBlocks.end())
+        continue;
+
+      ++I;
+      assert(NextBlock < getBlockCount());
+      ThisStream.push_back(NextBlock);
+    }
+    StreamMap.push_back(ThisStream);
+  }
+  return Error::success();
+}
+
 Error PDBFile::commit() {
   StreamWriter Writer(*Buffer);
 
@@ -371,5 +429,35 @@ Error PDBFile::commit() {
       return EC;
   }
 
+  if (Info) {
+    if (auto EC = Info->commit())
+      return EC;
+  }
+
+  if (Dbi) {
+    if (auto EC = Dbi->commit())
+      return EC;
+  }
+
+  if (Symbols) {
+    if (auto EC = Symbols->commit())
+      return EC;
+  }
+
+  if (Publics) {
+    if (auto EC = Publics->commit())
+      return EC;
+  }
+
+  if (Tpi) {
+    if (auto EC = Tpi->commit())
+      return EC;
+  }
+
+  if (Ipi) {
+    if (auto EC = Ipi->commit())
+      return EC;
+  }
+
   return Buffer->commit();
 }
\ No newline at end of file