1 files changed, 28 insertions, 13 deletions
diff --git a/llvm/lib/Object/ELFObjectFile.cpp b/llvm/lib/Object/ELFObjectFile.cpp
index 2cbfd8a..b9343d6 100644
--- a/llvm/lib/Object/ELFObjectFile.cpp
+++ b/llvm/lib/Object/ELFObjectFile.cpp
@@ -681,24 +681,39 @@ template <class ELFT>
 Expected<std::vector<BBAddrMap>> static readBBAddrMapImpl(
     const ELFFile<ELFT> &EF, std::optional<unsigned> TextSectionIndex) {
   using Elf_Shdr = typename ELFT::Shdr;
+  bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
   std::vector<BBAddrMap> BBAddrMaps;
+
   const auto &Sections = cantFail(EF.sections());
-  for (const Elf_Shdr &Sec : Sections) {
+  auto IsMatch = [&](const Elf_Shdr &Sec) -> Expected<bool> {
     if (Sec.sh_type != ELF::SHT_LLVM_BB_ADDR_MAP &&
         Sec.sh_type != ELF::SHT_LLVM_BB_ADDR_MAP_V0)
-      continue;
-    if (TextSectionIndex) {
-      Expected<const Elf_Shdr *> TextSecOrErr = EF.getSection(Sec.sh_link);
-      if (!TextSecOrErr)
-        return createError("unable to get the linked-to section for " +
-                           describe(EF, Sec) + ": " +
-                           toString(TextSecOrErr.takeError()));
-      if (*TextSectionIndex != std::distance(Sections.begin(), *TextSecOrErr))
-        continue;
-    }
-    Expected<std::vector<BBAddrMap>> BBAddrMapOrErr = EF.decodeBBAddrMap(Sec);
+      return false;
+    if (!TextSectionIndex)
+      return true;
+    Expected<const Elf_Shdr *> TextSecOrErr = EF.getSection(Sec.sh_link);
+    if (!TextSecOrErr)
+      return createError("unable to get the linked-to section for " +
+                         describe(EF, Sec) + ": " +
+                         toString(TextSecOrErr.takeError()));
+    if (*TextSectionIndex != std::distance(Sections.begin(), *TextSecOrErr))
+      return false;
+    return true;
+  };
+
+  Expected<MapVector<const Elf_Shdr *, const Elf_Shdr *>> SectionRelocMapOrErr =
+      EF.getSectionAndRelocations(IsMatch);
+  if (!SectionRelocMapOrErr)
+    return SectionRelocMapOrErr.takeError();
+
+  for (auto const &[Sec, RelocSec] : *SectionRelocMapOrErr) {
+    if (IsRelocatable && !RelocSec)
+      return createError("unable to get relocation section for " +
+                         describe(EF, *Sec));
+    Expected<std::vector<BBAddrMap>> BBAddrMapOrErr =
+        EF.decodeBBAddrMap(*Sec, RelocSec);
     if (!BBAddrMapOrErr)
-      return createError("unable to read " + describe(EF, Sec) + ": " +
+      return createError("unable to read " + describe(EF, *Sec) + ": " +
                          toString(BBAddrMapOrErr.takeError()));
     std::move(BBAddrMapOrErr->begin(), BBAddrMapOrErr->end(),
               std::back_inserter(BBAddrMaps));