1 files changed, 9 insertions, 12 deletions
diff --git a/llvm/tools/llvm-profgen/ProfiledBinary.cpp b/llvm/tools/llvm-profgen/ProfiledBinary.cpp
index 6865e36..94728ce 100644
--- a/llvm/tools/llvm-profgen/ProfiledBinary.cpp
+++ b/llvm/tools/llvm-profgen/ProfiledBinary.cpp
@@ -250,14 +250,12 @@ void ProfiledBinary::load() {
 
   DisassembleFunctionSet.insert_range(DisassembleFunctions);
 
-  if (auto *ELFObj = dyn_cast<ELFObjectFileBase>(Obj)) {
-    checkPseudoProbe(ELFObj);
-    if (UsePseudoProbes)
-      populateElfSymbolAddressList(ELFObj);
+  checkPseudoProbe(Obj);
+  if (UsePseudoProbes)
+    populateSymbolAddressList(Obj);
 
-    if (ShowDisassemblyOnly)
-      decodePseudoProbe(ELFObj);
-  }
+  if (ShowDisassemblyOnly)
+    decodePseudoProbe(Obj);
 
   // Disassemble the text sections.
   disassemble(Obj);
@@ -417,7 +415,7 @@ void ProfiledBinary::setPreferredTextSegmentAddresses(const ObjectFile *Obj) {
     llvm_unreachable("invalid object format");
 }
 
-void ProfiledBinary::checkPseudoProbe(const ELFObjectFileBase *Obj) {
+void ProfiledBinary::checkPseudoProbe(const ObjectFile *Obj) {
   if (UseDwarfCorrelation)
     return;
 
@@ -440,7 +438,7 @@ void ProfiledBinary::checkPseudoProbe(const ELFObjectFileBase *Obj) {
   UsePseudoProbes = HasProbeDescSection && HasPseudoProbeSection;
 }
 
-void ProfiledBinary::decodePseudoProbe(const ELFObjectFileBase *Obj) {
+void ProfiledBinary::decodePseudoProbe(const ObjectFile *Obj) {
   if (!UsePseudoProbes)
     return;
 
@@ -511,7 +509,7 @@ void ProfiledBinary::decodePseudoProbe(const ELFObjectFileBase *Obj) {
 void ProfiledBinary::decodePseudoProbe() {
   OwningBinary<Binary> OBinary = unwrapOrError(createBinary(Path), Path);
   Binary &ExeBinary = *OBinary.getBinary();
-  auto *Obj = cast<ELFObjectFileBase>(&ExeBinary);
+  auto *Obj = cast<ObjectFile>(&ExeBinary);
   decodePseudoProbe(Obj);
 }
 
@@ -809,8 +807,7 @@ void ProfiledBinary::checkUseFSDiscriminator(
   }
 }
 
-void ProfiledBinary::populateElfSymbolAddressList(
-    const ELFObjectFileBase *Obj) {
+void ProfiledBinary::populateSymbolAddressList(const ObjectFile *Obj) {
   // Create a mapping from virtual address to symbol GUID and the other way
   // around.
   StringRef FileName = Obj->getFileName();