1 files changed, 116 insertions, 7 deletions

diff --git a/lib/DebugInfo/PDB/DIA/DIASession.cpp b/lib/DebugInfo/PDB/DIA/DIASession.cpp
index b8aaebbf7380..d81f59400eb3 100644
--- a/lib/DebugInfo/PDB/DIA/DIASession.cpp
+++ b/lib/DebugInfo/PDB/DIA/DIASession.cpp
@@ -9,7 +9,9 @@
 #include "llvm/DebugInfo/PDB/DIA/DIASession.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/DebugInfo/PDB/DIA/DIAEnumDebugStreams.h"
+#include "llvm/DebugInfo/PDB/DIA/DIAEnumInjectedSources.h"
 #include "llvm/DebugInfo/PDB/DIA/DIAEnumLineNumbers.h"
+#include "llvm/DebugInfo/PDB/DIA/DIAEnumSectionContribs.h"
 #include "llvm/DebugInfo/PDB/DIA/DIAEnumSourceFiles.h"
 #include "llvm/DebugInfo/PDB/DIA/DIAEnumTables.h"
 #include "llvm/DebugInfo/PDB/DIA/DIAError.h"
@@ -104,7 +106,7 @@ Error DIASession::createFromPdb(StringRef Path,
   if (!llvm::convertUTF8ToUTF16String(Path, Path16))
     return make_error<GenericError>(generic_error_code::invalid_path);
 
-  const wchar_t *Path16Str = reinterpret_cast<const wchar_t*>(Path16.data());
+  const wchar_t *Path16Str = reinterpret_cast<const wchar_t *>(Path16.data());
   HRESULT HR;
   if (FAILED(HR = DiaDataSource->loadDataFromPdb(Path16Str))) {
     return ErrorFromHResult(HR, "Calling loadDataFromPdb {0}", Path);
@@ -148,8 +150,8 @@ uint64_t DIASession::getLoadAddress() const {
   return (success) ? LoadAddress : 0;
 }
 
-void DIASession::setLoadAddress(uint64_t Address) {
-  Session->put_loadAddress(Address);
+bool DIASession::setLoadAddress(uint64_t Address) {
+  return (S_OK == Session->put_loadAddress(Address));
 }
 
 std::unique_ptr<PDBSymbolExe> DIASession::getGlobalScope() {
@@ -164,6 +166,28 @@ std::unique_ptr<PDBSymbolExe> DIASession::getGlobalScope() {
   return ExeSymbol;
 }
 
+bool DIASession::addressForVA(uint64_t VA, uint32_t &Section,
+                              uint32_t &Offset) const {
+  DWORD ArgSection, ArgOffset = 0;
+  if (S_OK == Session->addressForVA(VA, &ArgSection, &ArgOffset)) {
+    Section = static_cast<uint32_t>(ArgSection);
+    Offset = static_cast<uint32_t>(ArgOffset);
+    return true;
+  }
+  return false;
+}
+
+bool DIASession::addressForRVA(uint32_t RVA, uint32_t &Section,
+                               uint32_t &Offset) const {
+  DWORD ArgSection, ArgOffset = 0;
+  if (S_OK == Session->addressForRVA(RVA, &ArgSection, &ArgOffset)) {
+    Section = static_cast<uint32_t>(ArgSection);
+    Offset = static_cast<uint32_t>(ArgOffset);
+    return true;
+  }
+  return false;
+}
+
 std::unique_ptr<PDBSymbol> DIASession::getSymbolById(uint32_t SymbolId) const {
   CComPtr<IDiaSymbol> LocatedSymbol;
   if (S_OK != Session->symbolById(SymbolId, &LocatedSymbol))
@@ -190,6 +214,31 @@ DIASession::findSymbolByAddress(uint64_t Address, PDB_SymType Type) const {
   return PDBSymbol::create(*this, std::move(RawSymbol));
 }
 
+std::unique_ptr<PDBSymbol> DIASession::findSymbolByRVA(uint32_t RVA,
+                                                       PDB_SymType Type) const {
+  enum SymTagEnum EnumVal = static_cast<enum SymTagEnum>(Type);
+
+  CComPtr<IDiaSymbol> Symbol;
+  if (S_OK != Session->findSymbolByRVA(RVA, EnumVal, &Symbol))
+    return nullptr;
+
+  auto RawSymbol = llvm::make_unique<DIARawSymbol>(*this, Symbol);
+  return PDBSymbol::create(*this, std::move(RawSymbol));
+}
+
+std::unique_ptr<PDBSymbol>
+DIASession::findSymbolBySectOffset(uint32_t Sect, uint32_t Offset,
+                                   PDB_SymType Type) const {
+  enum SymTagEnum EnumVal = static_cast<enum SymTagEnum>(Type);
+
+  CComPtr<IDiaSymbol> Symbol;
+  if (S_OK != Session->findSymbolByAddr(Sect, Offset, EnumVal, &Symbol))
+    return nullptr;
+
+  auto RawSymbol = llvm::make_unique<DIARawSymbol>(*this, Symbol);
+  return PDBSymbol::create(*this, std::move(RawSymbol));
+}
+
 std::unique_ptr<IPDBEnumLineNumbers>
 DIASession::findLineNumbers(const PDBSymbolCompiland &Compiland,
                             const IPDBSourceFile &File) const {
@@ -198,9 +247,8 @@ DIASession::findLineNumbers(const PDBSymbolCompiland &Compiland,
   const DIASourceFile &RawFile = static_cast<const DIASourceFile &>(File);
 
   CComPtr<IDiaEnumLineNumbers> LineNumbers;
-  if (S_OK !=
-      Session->findLines(RawCompiland.getDiaSymbol(), RawFile.getDiaFile(),
-                         &LineNumbers))
+  if (S_OK != Session->findLines(RawCompiland.getDiaSymbol(),
+                                 RawFile.getDiaFile(), &LineNumbers))
     return nullptr;
 
   return llvm::make_unique<DIAEnumLineNumbers>(LineNumbers);
@@ -209,7 +257,31 @@ DIASession::findLineNumbers(const PDBSymbolCompiland &Compiland,
 std::unique_ptr<IPDBEnumLineNumbers>
 DIASession::findLineNumbersByAddress(uint64_t Address, uint32_t Length) const {
   CComPtr<IDiaEnumLineNumbers> LineNumbers;
-  if (S_OK != Session->findLinesByVA(Address, Length, &LineNumbers))
+  if (S_OK != Session->findLinesByVA(Address, Length, &LineNumbers)) {
+    ULONGLONG LoadAddr = 0;
+    if (S_OK != Session->get_loadAddress(&LoadAddr))
+      return nullptr;
+    DWORD RVA = static_cast<DWORD>(Address - LoadAddr);
+    if (S_OK != Session->findLinesByRVA(RVA, Length, &LineNumbers))
+      return nullptr;
+  }
+  return llvm::make_unique<DIAEnumLineNumbers>(LineNumbers);
+}
+
+std::unique_ptr<IPDBEnumLineNumbers>
+DIASession::findLineNumbersByRVA(uint32_t RVA, uint32_t Length) const {
+  CComPtr<IDiaEnumLineNumbers> LineNumbers;
+  if (S_OK != Session->findLinesByRVA(RVA, Length, &LineNumbers))
+    return nullptr;
+
+  return llvm::make_unique<DIAEnumLineNumbers>(LineNumbers);
+}
+
+std::unique_ptr<IPDBEnumLineNumbers>
+DIASession::findLineNumbersBySectOffset(uint32_t Section, uint32_t Offset,
+                                        uint32_t Length) const {
+  CComPtr<IDiaEnumLineNumbers> LineNumbers;
+  if (S_OK != Session->findLinesByAddr(Section, Offset, Length, &LineNumbers))
     return nullptr;
 
   return llvm::make_unique<DIAEnumLineNumbers>(LineNumbers);
@@ -310,3 +382,40 @@ std::unique_ptr<IPDBEnumTables> DIASession::getEnumTables() const {
 
   return llvm::make_unique<DIAEnumTables>(DiaEnumerator);
 }
+
+template <class T> static CComPtr<T> getTableEnumerator(IDiaSession &Session) {
+  CComPtr<T> Enumerator;
+  CComPtr<IDiaEnumTables> ET;
+  CComPtr<IDiaTable> Table;
+  ULONG Count = 0;
+
+  if (Session.getEnumTables(&ET) != S_OK)
+    return nullptr;
+
+  while (ET->Next(1, &Table, &Count) == S_OK && Count == 1) {
+    // There is only one table that matches the given iid
+    if (S_OK == Table->QueryInterface(__uuidof(T), (void **)&Enumerator))
+      break;
+    Table.Release();
+  }
+  return Enumerator;
+}
+std::unique_ptr<IPDBEnumInjectedSources>
+DIASession::getInjectedSources() const {
+  CComPtr<IDiaEnumInjectedSources> Files =
+      getTableEnumerator<IDiaEnumInjectedSources>(*Session);
+  if (!Files)
+    return nullptr;
+
+  return llvm::make_unique<DIAEnumInjectedSources>(*this, Files);
+}
+
+std::unique_ptr<IPDBEnumSectionContribs>
+DIASession::getSectionContribs() const {
+  CComPtr<IDiaEnumSectionContribs> Sections =
+      getTableEnumerator<IDiaEnumSectionContribs>(*Session);
+  if (!Sections)
+    return nullptr;
+
+  return llvm::make_unique<DIAEnumSectionContribs>(*this, Sections);
+}