2 files changed, 15 insertions, 27 deletions
diff --git a/llvm/lib/DebugInfo/PDB/PDBContext.cpp b/llvm/lib/DebugInfo/PDB/PDBContext.cpp
index 83f27c7fa3d..ca2ae6665ce 100644
--- a/llvm/lib/DebugInfo/PDB/PDBContext.cpp
+++ b/llvm/lib/DebugInfo/PDB/PDBContext.cpp
@@ -21,24 +21,11 @@ using namespace llvm;
 using namespace llvm::object;
 
 PDBContext::PDBContext(const COFFObjectFile &Object,
-                       std::unique_ptr<IPDBSession> PDBSession,
-                       bool RelativeAddress)
+                       std::unique_ptr<IPDBSession> PDBSession)
     : DIContext(CK_PDB), Session(std::move(PDBSession)) {
-  if (!RelativeAddress) {
-    uint64_t ImageBase = 0;
-    if (Object.is64()) {
-      const pe32plus_header *Header = nullptr;
-      Object.getPE32PlusHeader(Header);
-      if (Header)
-        ImageBase = Header->ImageBase;
-    } else {
-      const pe32_header *Header = nullptr;
-      Object.getPE32Header(Header);
-      if (Header)
-        ImageBase = static_cast<uint64_t>(Header->ImageBase);
-    }
-    Session->setLoadAddress(ImageBase);
-  }
+  ErrorOr<uint64_t> ImageBase = Object.getImageBase();
+  if (ImageBase)
+    Session->setLoadAddress(ImageBase.get());
 }
 
 void PDBContext::dump(raw_ostream &OS, DIDumpType DumpType) {}
diff --git a/llvm/lib/Object/COFFObjectFile.cpp b/llvm/lib/Object/COFFObjectFile.cpp
index d3f604a8d35..efb3ea04083 100644
--- a/llvm/lib/Object/COFFObjectFile.cpp
+++ b/llvm/lib/Object/COFFObjectFile.cpp
@@ -174,10 +174,7 @@ ErrorOr<uint64_t> COFFObjectFile::getSymbolAddress(DataRefImpl Ref) const {
 
   // The section VirtualAddress does not include ImageBase, and we want to
   // return virtual addresses.
-  if (PE32Header)
-    Result += PE32Header->ImageBase;
-  else if (PE32PlusHeader)
-    Result += PE32PlusHeader->ImageBase;
+  Result += getImageBase().get();
 
   return Result;
 }
@@ -274,10 +271,7 @@ uint64_t COFFObjectFile::getSectionAddress(DataRefImpl Ref) const {
 
   // The section VirtualAddress does not include ImageBase, and we want to
   // return virtual addresses.
-  if (PE32Header)
-    Result += PE32Header->ImageBase;
-  else if (PE32PlusHeader)
-    Result += PE32PlusHeader->ImageBase;
+  Result += getImageBase().get();
   return Result;
 }
 
@@ -424,10 +418,17 @@ std::error_code COFFObjectFile::initSymbolTablePtr() {
   return std::error_code();
 }
 
+ErrorOr<uint64_t> COFFObjectFile::getImageBase() const {
+  if (PE32Header)
+    return uint64_t(PE32Header->ImageBase);
+  else if (PE32PlusHeader)
+    return uint64_t(PE32PlusHeader->ImageBase);
+  return object_error::parse_failed;
+}
+
 // Returns the file offset for the given VA.
 std::error_code COFFObjectFile::getVaPtr(uint64_t Addr, uintptr_t &Res) const {
-  uint64_t ImageBase = PE32Header ? (uint64_t)PE32Header->ImageBase
-                                  : (uint64_t)PE32PlusHeader->ImageBase;
+  uint64_t ImageBase = getImageBase().get();
   uint64_t Rva = Addr - ImageBase;
   assert(Rva <= UINT32_MAX);
   return getRvaPtr((uint32_t)Rva, Res);