Resubmit "Support embedding natvis files in PDBs."

The issue causing this to fail in certain configurations
should be fixed.

It was due to the fact that DIA apparently expects there to be
a null string at ID 1 in the string table.  I'm not sure why this
is important but it seems to make a difference, so set it.

llvm-svn: 328002

3 files changed, 161 insertions, 13 deletions

diff --git a/llvm/lib/DebugInfo/PDB/Native/InfoStreamBuilder.cpp b/llvm/lib/DebugInfo/PDB/Native/InfoStreamBuilder.cpp
index a20b45111cf..54d6835f112 100644
--- a/llvm/lib/DebugInfo/PDB/Native/InfoStreamBuilder.cpp
+++ b/llvm/lib/DebugInfo/PDB/Native/InfoStreamBuilder.cpp
@@ -73,5 +73,6 @@ Error InfoStreamBuilder::commit(const msf::MSFLayout &Layout,
     if (auto EC = Writer.writeEnum(E))
       return EC;
   }
+  assert(Writer.bytesRemaining() == 0);
   return Error::success();
 }
diff --git a/llvm/lib/DebugInfo/PDB/Native/PDBFileBuilder.cpp b/llvm/lib/DebugInfo/PDB/Native/PDBFileBuilder.cpp
index 1cb890ff799..03d9f3d7bed 100644
--- a/llvm/lib/DebugInfo/PDB/Native/PDBFileBuilder.cpp
+++ b/llvm/lib/DebugInfo/PDB/Native/PDBFileBuilder.cpp
@@ -24,6 +24,8 @@
 #include "llvm/DebugInfo/PDB/Native/TpiStreamBuilder.h"
 #include "llvm/Support/BinaryStream.h"
 #include "llvm/Support/BinaryStreamWriter.h"
+#include "llvm/Support/JamCRC.h"
+#include "llvm/Support/Path.h"
 
 using namespace llvm;
 using namespace llvm::codeview;
@@ -32,7 +34,8 @@ using namespace llvm::pdb;
 using namespace llvm::support;
 
 PDBFileBuilder::PDBFileBuilder(BumpPtrAllocator &Allocator)
-    : Allocator(Allocator) {}
+    : Allocator(Allocator), InjectedSourceHashTraits(Strings),
+      InjectedSourceTable(2, InjectedSourceHashTraits) {}
 
 PDBFileBuilder::~PDBFileBuilder() {}
 
@@ -80,14 +83,45 @@ GSIStreamBuilder &PDBFileBuilder::getGsiBuilder() {
   return *Gsi;
 }
 
-Error PDBFileBuilder::addNamedStream(StringRef Name, uint32_t Size) {
+Expected<uint32_t> PDBFileBuilder::allocateNamedStream(StringRef Name,
+                                                       uint32_t Size) {
   auto ExpectedStream = Msf->addStream(Size);
-  if (!ExpectedStream)
-    return ExpectedStream.takeError();
-  NamedStreams.set(Name, *ExpectedStream);
+  if (ExpectedStream)
+    NamedStreams.set(Name, *ExpectedStream);
+  return ExpectedStream;
+}
+
+Error PDBFileBuilder::addNamedStream(StringRef Name, StringRef Data) {
+  Expected<uint32_t> ExpectedIndex = allocateNamedStream(Name, Data.size());
+  if (!ExpectedIndex)
+    return ExpectedIndex.takeError();
+  assert(NamedStreamData.count(*ExpectedIndex) == 0);
+  NamedStreamData[*ExpectedIndex] = Data;
   return Error::success();
 }
 
+void PDBFileBuilder::addInjectedSource(StringRef Name,
+                                       std::unique_ptr<MemoryBuffer> Buffer) {
+  // Stream names must be exact matches, since they get looked up in a hash
+  // table and the hash value is dependent on the exact contents of the string.
+  // link.exe lowercases a path and converts / to \, so we must do the same.
+  SmallString<64> VName;
+  sys::path::native(Name.lower(), VName);
+
+  uint32_t NI = getStringTableBuilder().insert(Name);
+  uint32_t VNI = getStringTableBuilder().insert(VName);
+
+  InjectedSourceDescriptor Desc;
+  Desc.Content = std::move(Buffer);
+  Desc.NameIndex = NI;
+  Desc.VNameIndex = VNI;
+  Desc.StreamName = "/src/files/";
+
+  Desc.StreamName += VName;
+
+  InjectedSources.push_back(std::move(Desc));
+}
+
 Expected<msf::MSFLayout> PDBFileBuilder::finalizeMsfLayout() {
 
   if (Ipi && Ipi->getRecordCount() > 0) {
@@ -101,15 +135,13 @@ Expected<msf::MSFLayout> PDBFileBuilder::finalizeMsfLayout() {
 
   uint32_t StringsLen = Strings.calculateSerializedSize();
 
-  if (auto EC = addNamedStream("/names", StringsLen))
-    return std::move(EC);
-  if (auto EC = addNamedStream("/LinkInfo", 0))
-    return std::move(EC);
+  Expected<uint32_t> SN = allocateNamedStream("/names", StringsLen);
+  if (!SN)
+    return SN.takeError();
+  SN = allocateNamedStream("/LinkInfo", 0);
+  if (!SN)
+    return SN.takeError();
 
-  if (Info) {
-    if (auto EC = Info->finalizeMsfLayout())
-      return std::move(EC);
-  }
   if (Dbi) {
     if (auto EC = Dbi->finalizeMsfLayout())
       return std::move(EC);
@@ -132,6 +164,46 @@ Expected<msf::MSFLayout> PDBFileBuilder::finalizeMsfLayout() {
     }
   }
 
+  if (!InjectedSources.empty()) {
+    for (const auto &IS : InjectedSources) {
+      JamCRC CRC(0);
+      CRC.update(makeArrayRef(IS.Content->getBufferStart(),
+                              IS.Content->getBufferSize()));
+
+      SrcHeaderBlockEntry Entry;
+      ::memset(&Entry, 0, sizeof(SrcHeaderBlockEntry));
+      Entry.Size = sizeof(SrcHeaderBlockEntry);
+      Entry.FileSize = IS.Content->getBufferSize();
+      Entry.FileNI = IS.NameIndex;
+      Entry.VFileNI = IS.VNameIndex;
+      Entry.IsVirtual = 0;
+      Entry.Version =
+          static_cast<uint32_t>(PdbRaw_SrcHeaderBlockVer::SrcVerOne);
+      Entry.CRC = CRC.getCRC();
+      StringRef VName = getStringTableBuilder().getStringForId(IS.VNameIndex);
+      InjectedSourceTable.set_as(VName, std::move(Entry));
+    }
+
+    uint32_t SrcHeaderBlockSize =
+        sizeof(SrcHeaderBlockHeader) +
+        InjectedSourceTable.calculateSerializedLength();
+    SN = allocateNamedStream("/src/headerblock", SrcHeaderBlockSize);
+    if (!SN)
+      return SN.takeError();
+    for (const auto &IS : InjectedSources) {
+      SN = allocateNamedStream(IS.StreamName, IS.Content->getBufferSize());
+      if (!SN)
+        return SN.takeError();
+    }
+  }
+
+  // Do this last, since it relies on the named stream map being complete, and
+  // that can be updated by previous steps in the finalization.
+  if (Info) {
+    if (auto EC = Info->finalizeMsfLayout())
+      return std::move(EC);
+  }
+
   return Msf->build();
 }
 
@@ -167,6 +239,45 @@ void PDBFileBuilder::commitFpm(WritableBinaryStream &MsfBuffer,
   assert(FpmWriter.bytesRemaining() == 0);
 }
 
+void PDBFileBuilder::commitSrcHeaderBlock(WritableBinaryStream &MsfBuffer,
+                                          const msf::MSFLayout &Layout) {
+  assert(!InjectedSourceTable.empty());
+
+  uint32_t SN = cantFail(getNamedStreamIndex("/src/headerblock"));
+  auto Stream = WritableMappedBlockStream::createIndexedStream(
+      Layout, MsfBuffer, SN, Allocator);
+  BinaryStreamWriter Writer(*Stream);
+
+  SrcHeaderBlockHeader Header;
+  ::memset(&Header, 0, sizeof(Header));
+  Header.Version = static_cast<uint32_t>(PdbRaw_SrcHeaderBlockVer::SrcVerOne);
+  Header.Size = Writer.bytesRemaining();
+
+  cantFail(Writer.writeObject(Header));
+  cantFail(InjectedSourceTable.commit(Writer));
+
+  assert(Writer.bytesRemaining() == 0);
+}
+
+void PDBFileBuilder::commitInjectedSources(WritableBinaryStream &MsfBuffer,
+                                           const msf::MSFLayout &Layout) {
+  if (InjectedSourceTable.empty())
+    return;
+
+  commitSrcHeaderBlock(MsfBuffer, Layout);
+
+  for (const auto &IS : InjectedSources) {
+    uint32_t SN = cantFail(getNamedStreamIndex(IS.StreamName));
+
+    auto SourceStream = WritableMappedBlockStream::createIndexedStream(
+        Layout, MsfBuffer, SN, Allocator);
+    BinaryStreamWriter SourceWriter(*SourceStream);
+    assert(SourceWriter.bytesRemaining() == IS.Content->getBufferSize());
+    cantFail(SourceWriter.writeBytes(
+        arrayRefFromStringRef(IS.Content->getBuffer())));
+  }
+}
+
 Error PDBFileBuilder::commit(StringRef Filename) {
   assert(!Filename.empty());
   auto ExpectedLayout = finalizeMsfLayout();
@@ -219,6 +330,17 @@ Error PDBFileBuilder::commit(StringRef Filename) {
   if (auto EC = Strings.commit(NSWriter))
     return EC;
 
+  for (const auto &NSE : NamedStreamData) {
+    if (NSE.second.empty())
+      continue;
+
+    auto NS = WritableMappedBlockStream::createIndexedStream(
+        Layout, Buffer, NSE.first, Allocator);
+    BinaryStreamWriter NSW(*NS);
+    if (auto EC = NSW.writeBytes(arrayRefFromStringRef(NSE.second)))
+      return EC;
+  }
+
   if (Info) {
     if (auto EC = Info->commit(Layout, Buffer))
       return EC;
@@ -251,6 +373,8 @@ Error PDBFileBuilder::commit(StringRef Filename) {
   InfoStreamHeader *H = reinterpret_cast<InfoStreamHeader *>(
       FOB->getBufferStart() + InfoStreamFileOffset);
 
+  commitInjectedSources(Buffer, Layout);
+
   // Set the build id at the very end, after every other byte of the PDB
   // has been written.
   // FIXME: Use a hash of the PDB rather than time(nullptr) for the signature.
diff --git a/llvm/lib/DebugInfo/PDB/Native/PDBStringTableBuilder.cpp b/llvm/lib/DebugInfo/PDB/Native/PDBStringTableBuilder.cpp
index ece3e00b1a8..13106ccbc32 100644
--- a/llvm/lib/DebugInfo/PDB/Native/PDBStringTableBuilder.cpp
+++ b/llvm/lib/DebugInfo/PDB/Native/PDBStringTableBuilder.cpp
@@ -21,10 +21,33 @@ using namespace llvm::support;
 using namespace llvm::support::endian;
 using namespace llvm::pdb;
 
+StringTableHashTraits::StringTableHashTraits(PDBStringTableBuilder &Table)
+    : Table(&Table) {}
+
+uint32_t StringTableHashTraits::hashLookupKey(StringRef S) const {
+  return Table->getIdForString(S);
+}
+
+StringRef StringTableHashTraits::storageKeyToLookupKey(uint32_t Offset) const {
+  return Table->getStringForId(Offset);
+}
+
+uint32_t StringTableHashTraits::lookupKeyToStorageKey(StringRef S) {
+  return Table->insert(S);
+}
+
 uint32_t PDBStringTableBuilder::insert(StringRef S) {
   return Strings.insert(S);
 }
 
+uint32_t PDBStringTableBuilder::getIdForString(StringRef S) const {
+  return Strings.getIdForString(S);
+}
+
+StringRef PDBStringTableBuilder::getStringForId(uint32_t Id) const {
+  return Strings.getStringForId(Id);
+}
+
 static uint32_t computeBucketCount(uint32_t NumStrings) {
   // The /names stream is basically an on-disk open-addressing hash table.
   // Hash collisions are resolved by linear probing. We cannot make