Support embedding natvis files in PDBs.

Natvis is a debug language supported by Visual Studio for
specifying custom visualizers.  The /NATVIS option is an
undocumented link.exe flag which will take a .natvis file
and "inject" it into the PDB.  This way, you can ship the
debug visualizers for a program along with the PDB, which
is very useful for postmortem debugging.

This is implemented by adding a new "named stream" to the
PDB with a special name of /src/files/<natvis file name>
and simply copying the contents of the xml into this file.

Additionally, we need to emit a single stream named
/src/headerblock which contains a hash table of embedded
files to records describing them.

This patch adds this functionality, including the /NATVIS
option to lld-link.

Differential Revision: https://reviews.llvm.org/D44328

llvm-svn: 327895

6 files changed, 182 insertions, 25 deletions

diff --git a/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp b/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp
index ccc20eb7488..0f155a95d60 100644
--- a/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp
+++ b/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp
@@ -109,7 +109,7 @@ Error DebugChecksumsSubsection::commit(BinaryStreamWriter &Writer) const {
 }
 
 uint32_t DebugChecksumsSubsection::mapChecksumOffset(StringRef FileName) const {
-  uint32_t Offset = Strings.getStringId(FileName);
+  uint32_t Offset = Strings.getIdForString(FileName);
   auto Iter = OffsetMap.find(Offset);
   assert(Iter != OffsetMap.end());
   return Iter->second;
diff --git a/llvm/lib/DebugInfo/CodeView/DebugCrossImpSubsection.cpp b/llvm/lib/DebugInfo/CodeView/DebugCrossImpSubsection.cpp
index 88c0076915b..9a3d3e3e247 100644
--- a/llvm/lib/DebugInfo/CodeView/DebugCrossImpSubsection.cpp
+++ b/llvm/lib/DebugInfo/CodeView/DebugCrossImpSubsection.cpp
@@ -80,13 +80,13 @@ Error DebugCrossModuleImportsSubsection::commit(
     Ids.push_back(&M);
 
   std::sort(Ids.begin(), Ids.end(), [this](const T &L1, const T &L2) {
-    return Strings.getStringId(L1->getKey()) <
-           Strings.getStringId(L2->getKey());
+    return Strings.getIdForString(L1->getKey()) <
+           Strings.getIdForString(L2->getKey());
   });
 
   for (const auto &Item : Ids) {
     CrossModuleImport Imp;
-    Imp.ModuleNameOffset = Strings.getStringId(Item->getKey());
+    Imp.ModuleNameOffset = Strings.getIdForString(Item->getKey());
     Imp.Count = Item->getValue().size();
     if (auto EC = Writer.writeObject(Imp))
       return EC;
diff --git a/llvm/lib/DebugInfo/CodeView/DebugStringTableSubsection.cpp b/llvm/lib/DebugInfo/CodeView/DebugStringTableSubsection.cpp
index d723282eb71..c731b68625c 100644
--- a/llvm/lib/DebugInfo/CodeView/DebugStringTableSubsection.cpp
+++ b/llvm/lib/DebugInfo/CodeView/DebugStringTableSubsection.cpp
@@ -46,12 +46,15 @@ DebugStringTableSubsection::DebugStringTableSubsection()
     : DebugSubsection(DebugSubsectionKind::StringTable) {}
 
 uint32_t DebugStringTableSubsection::insert(StringRef S) {
-  auto P = Strings.insert({S, StringSize});
+  auto P = StringToId.insert({S, StringSize});
 
   // If a given string didn't exist in the string table, we want to increment
-  // the string table size.
-  if (P.second)
+  // the string table size and insert it into the reverse lookup.
+  if (P.second) {
+    IdToString.insert({P.first->getValue(), P.first->getKey()});
     StringSize += S.size() + 1; // +1 for '\0'
+  }
+
   return P.first->second;
 }
 
@@ -67,7 +70,7 @@ Error DebugStringTableSubsection::commit(BinaryStreamWriter &Writer) const {
   if (auto EC = Writer.writeCString(StringRef()))
     return EC;
 
-  for (auto &Pair : Strings) {
+  for (auto &Pair : StringToId) {
     StringRef S = Pair.getKey();
     uint32_t Offset = Begin + Pair.getValue();
     Writer.setOffset(Offset);
@@ -81,10 +84,16 @@ Error DebugStringTableSubsection::commit(BinaryStreamWriter &Writer) const {
   return Error::success();
 }
 
-uint32_t DebugStringTableSubsection::size() const { return Strings.size(); }
+uint32_t DebugStringTableSubsection::size() const { return StringToId.size(); }
+
+uint32_t DebugStringTableSubsection::getIdForString(StringRef S) const {
+  auto Iter = StringToId.find(S);
+  assert(Iter != StringToId.end());
+  return Iter->second;
+}
 
-uint32_t DebugStringTableSubsection::getStringId(StringRef S) const {
-  auto Iter = Strings.find(S);
-  assert(Iter != Strings.end());
+StringRef DebugStringTableSubsection::getStringForId(uint32_t Id) const {
+  auto Iter = IdToString.find(Id);
+  assert(Iter != IdToString.end());
   return Iter->second;
 }
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