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

8 files changed, 186 insertions, 32 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/DIA/DIAEnumInjectedSources.cpp b/llvm/lib/DebugInfo/PDB/DIA/DIAEnumInjectedSources.cpp
index 873f4a0a8db..d7c908e0459 100644
--- a/llvm/lib/DebugInfo/PDB/DIA/DIAEnumInjectedSources.cpp
+++ b/llvm/lib/DebugInfo/PDB/DIA/DIAEnumInjectedSources.cpp
@@ -30,8 +30,7 @@ DIAEnumInjectedSources::getChildAtIndex(uint32_t Index) const {
   if (S_OK != Enumerator->Item(Index, &Item))
     return nullptr;
 
-  return std::unique_ptr<IPDBInjectedSource>(
-      new DIAInjectedSource(Session, Item));
+  return std::unique_ptr<IPDBInjectedSource>(new DIAInjectedSource(Item));
 }
 
 std::unique_ptr<IPDBInjectedSource> DIAEnumInjectedSources::getNext() {
@@ -40,8 +39,7 @@ std::unique_ptr<IPDBInjectedSource> DIAEnumInjectedSources::getNext() {
   if (S_OK != Enumerator->Next(1, &Item, &NumFetched))
     return nullptr;
 
-  return std::unique_ptr<IPDBInjectedSource>(
-      new DIAInjectedSource(Session, Item));
+  return std::unique_ptr<IPDBInjectedSource>(new DIAInjectedSource(Item));
 }
 
 void DIAEnumInjectedSources::reset() { Enumerator->Reset(); }
diff --git a/llvm/lib/DebugInfo/PDB/DIA/DIAInjectedSource.cpp b/llvm/lib/DebugInfo/PDB/DIA/DIAInjectedSource.cpp
index fbd169524b6..1d642f221d7 100644
--- a/llvm/lib/DebugInfo/PDB/DIA/DIAInjectedSource.cpp
+++ b/llvm/lib/DebugInfo/PDB/DIA/DIAInjectedSource.cpp
@@ -16,9 +16,8 @@
 using namespace llvm;
 using namespace llvm::pdb;
 
-DIAInjectedSource::DIAInjectedSource(const DIASession &Session,
-                                     CComPtr<IDiaInjectedSource> DiaSourceFile)
-    : Session(Session), SourceFile(DiaSourceFile) {}
+DIAInjectedSource::DIAInjectedSource(CComPtr<IDiaInjectedSource> DiaSourceFile)
+    : SourceFile(DiaSourceFile) {}
 
 uint32_t DIAInjectedSource::getCrc32() const {
   DWORD Crc;
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