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
diff --git a/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp b/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp
index ccc20eb..0f155a9 100644
--- a/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp
+++ b/llvm/lib/DebugInfo/CodeView/DebugChecksumsSubsection.cpp
@@ -109,7 +109,7 @@
 }
 
 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 88c0076..9a3d3e3 100644
--- a/llvm/lib/DebugInfo/CodeView/DebugCrossImpSubsection.cpp
+++ b/llvm/lib/DebugInfo/CodeView/DebugCrossImpSubsection.cpp
@@ -80,13 +80,13 @@
     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 d723282..c731b68 100644
--- a/llvm/lib/DebugInfo/CodeView/DebugStringTableSubsection.cpp
+++ b/llvm/lib/DebugInfo/CodeView/DebugStringTableSubsection.cpp
@@ -46,12 +46,15 @@
     : 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 @@
   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 @@
   return Error::success();
 }
 
-uint32_t DebugStringTableSubsection::size() const { return Strings.size(); }
+uint32_t DebugStringTableSubsection::size() const { return StringToId.size(); }
 
-uint32_t DebugStringTableSubsection::getStringId(StringRef S) const {
-  auto Iter = Strings.find(S);
-  assert(Iter != Strings.end());
+uint32_t DebugStringTableSubsection::getIdForString(StringRef S) const {
+  auto Iter = StringToId.find(S);
+  assert(Iter != StringToId.end());
+  return Iter->second;
+}
+
+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 873f4a0..d7c908e 100644
--- a/llvm/lib/DebugInfo/PDB/DIA/DIAEnumInjectedSources.cpp
+++ b/llvm/lib/DebugInfo/PDB/DIA/DIAEnumInjectedSources.cpp
@@ -30,8 +30,7 @@
   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 @@
   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 fbd1695..1d642f2 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 a20b451..54d6835 100644
--- a/llvm/lib/DebugInfo/PDB/Native/InfoStreamBuilder.cpp
+++ b/llvm/lib/DebugInfo/PDB/Native/InfoStreamBuilder.cpp
@@ -73,5 +73,6 @@
     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 1cb890f..03d9f3d 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::support;
 
 PDBFileBuilder::PDBFileBuilder(BumpPtrAllocator &Allocator)
-    : Allocator(Allocator) {}
+    : Allocator(Allocator), InjectedSourceHashTraits(Strings),
+      InjectedSourceTable(2, InjectedSourceHashTraits) {}
 
 PDBFileBuilder::~PDBFileBuilder() {}
 
@@ -80,14 +83,45 @@
   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 @@
 
   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 @@
     }
   }
 
+  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 @@
   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 @@
   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 @@
   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 ece3e00..13106cc 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::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