[llvm] [Object] Extract format-agnostic BBAddrMap decoder (PR #188435)

Haohai Wen via llvm-commits llvm-commits at lists.llvm.org
Mon Mar 30 19:15:26 PDT 2026


================
@@ -724,294 +724,125 @@ ELFFile<ELFT>::toMappedAddr(uint64_t VAddr, WarningHandler WarnHandler) const {
   return base() + Offset;
 }
 
-// Helper to extract and decode the next ULEB128 value as unsigned int.
-// Returns zero and sets ULEBSizeErr if the ULEB128 value exceeds the unsigned
-// int limit.
-// Also returns zero if ULEBSizeErr is already in an error state.
-// ULEBSizeErr is an out variable if an error occurs.
-template <typename IntTy, std::enable_if_t<std::is_unsigned_v<IntTy>, int> = 0>
-static IntTy readULEB128As(DataExtractor &Data, DataExtractor::Cursor &Cur,
-                           Error &ULEBSizeErr) {
-  // Bail out and do not extract data if ULEBSizeErr is already set.
-  if (ULEBSizeErr)
-    return 0;
-  uint64_t Offset = Cur.tell();
-  uint64_t Value = Data.getULEB128(Cur);
-  if (Value > std::numeric_limits<IntTy>::max()) {
-    ULEBSizeErr = createError("ULEB128 value at offset 0x" +
-                              Twine::utohexstr(Offset) + " exceeds UINT" +
-                              Twine(std::numeric_limits<IntTy>::digits) +
-                              "_MAX (0x" + Twine::utohexstr(Value) + ")");
-    return 0;
-  }
-  return static_cast<IntTy>(Value);
-}
-
-template <typename ELFT>
-static Expected<std::vector<BBAddrMap>>
-decodeBBAddrMapImpl(const ELFFile<ELFT> &EF,
-                    const typename ELFFile<ELFT>::Elf_Shdr &Sec,
-                    const typename ELFFile<ELFT>::Elf_Shdr *RelaSec,
-                    std::vector<PGOAnalysisMap> *PGOAnalyses) {
-  bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
-
-  // This DenseMap maps the offset of each function (the location of the
-  // reference to the function in the SHT_LLVM_BB_ADDR_MAP section) to the
-  // addend (the location of the function in the text section).
+/// Address extractor for ELF BB address map sections.
+class ELFBBAddrMapAddressExtractor : public AddressExtractor {
+  bool IsRelocatable;
+  // Maps the offset of each address field in the BB addr map section to the
+  // resolved function address (the relocation addend).
   llvm::DenseMap<uint64_t, uint64_t> FunctionOffsetTranslations;
-  if (IsRelocatable && RelaSec) {
-    assert(RelaSec &&
-           "Can't read a SHT_LLVM_BB_ADDR_MAP section in a relocatable "
-           "object file without providing a relocation section.");
-    if (RelaSec->sh_type == ELF::SHT_CREL) {
-      Expected<typename ELFFile<ELFT>::RelsOrRelas> Relas = EF.crels(*RelaSec);
-      if (!Relas)
-        return createError("unable to read CREL relocations for section " +
-                           describe(EF, Sec) + ": " +
-                           toString(Relas.takeError()));
-      for (typename ELFFile<ELFT>::Elf_Rela Rela : std::get<1>(*Relas)) {
-        FunctionOffsetTranslations[Rela.r_offset] = Rela.r_addend;
+  std::function<std::string()> DescribeSection;
+  std::unique_ptr<uint8_t[]> DecompressedContent;
+
+  ELFBBAddrMapAddressExtractor(
+      DataExtractor Data, bool IsRelocatable,
+      llvm::DenseMap<uint64_t, uint64_t> FunctionOffsetTranslations,
+      std::function<std::string()> DescribeSection,
+      std::unique_ptr<uint8_t[]> DecompressedContent)
+      : AddressExtractor(std::move(Data)), IsRelocatable(IsRelocatable),
+        FunctionOffsetTranslations(std::move(FunctionOffsetTranslations)),
+        DescribeSection(std::move(DescribeSection)),
+        DecompressedContent(std::move(DecompressedContent)) {}
+
+public:
+  template <typename ELFT>
+  static Expected<ELFBBAddrMapAddressExtractor>
+  create(const ELFFile<ELFT> &EF, const typename ELFFile<ELFT>::Elf_Shdr &Sec,
+         const typename ELFFile<ELFT>::Elf_Shdr *RelaSec) {
+    bool IsRelocatable = EF.getHeader().e_type == ELF::ET_REL;
+
+    // Build relocation offset-to-addend map.
+    llvm::DenseMap<uint64_t, uint64_t> FunctionOffsetTranslations;
+    if (IsRelocatable && RelaSec) {
+      assert(RelaSec &&
+             "Can't read a SHT_LLVM_BB_ADDR_MAP section in a relocatable "
+             "object file without providing a relocation section.");
+      if (RelaSec->sh_type == ELF::SHT_CREL) {
+        Expected<typename ELFFile<ELFT>::RelsOrRelas> Relas =
+            EF.crels(*RelaSec);
+        if (!Relas)
+          return createError("unable to read CREL relocations for section " +
+                             describe(EF, Sec) + ": " +
+                             toString(Relas.takeError()));
+        for (typename ELFFile<ELFT>::Elf_Rela Rela : std::get<1>(*Relas))
+          FunctionOffsetTranslations[Rela.r_offset] = Rela.r_addend;
+      } else {
+        Expected<typename ELFFile<ELFT>::Elf_Rela_Range> Relas =
+            EF.relas(*RelaSec);
+        if (!Relas)
+          return createError("unable to read relocations for section " +
+                             describe(EF, Sec) + ": " +
+                             toString(Relas.takeError()));
+        for (typename ELFFile<ELFT>::Elf_Rela Rela : *Relas)
+          FunctionOffsetTranslations[Rela.r_offset] = Rela.r_addend;
       }
-    } else {
-      Expected<typename ELFFile<ELFT>::Elf_Rela_Range> Relas =
-          EF.relas(*RelaSec);
-      if (!Relas)
-        return createError("unable to read relocations for section " +
-                           describe(EF, Sec) + ": " +
-                           toString(Relas.takeError()));
-      for (typename ELFFile<ELFT>::Elf_Rela Rela : *Relas)
-        FunctionOffsetTranslations[Rela.r_offset] = Rela.r_addend;
     }
-  }
-  auto GetAddressForRelocation =
-      [&](uint64_t RelocationOffsetInSection) -> Expected<uint64_t> {
-    auto FOTIterator =
-        FunctionOffsetTranslations.find(RelocationOffsetInSection);
-    if (FOTIterator == FunctionOffsetTranslations.end()) {
-      return createError("failed to get relocation data for offset: " +
-                         Twine::utohexstr(RelocationOffsetInSection) +
-                         " in section " + describe(EF, Sec));
-    }
-    return FOTIterator->second;
-  };
-  Expected<ArrayRef<uint8_t>> ContentsOrErr = EF.getSectionContents(Sec);
-  if (!ContentsOrErr)
-    return ContentsOrErr.takeError();
-  ArrayRef<uint8_t> Content = *ContentsOrErr;
 
-  // Decompress the section if needed.
-  std::unique_ptr<uint8_t[]> DecompressedContent;
-  if (Sec.sh_flags & llvm::ELF::SHF_COMPRESSED) {
-    Expected<StringRef> SectionNameOrErr = EF.getSectionName(Sec);
-    if (!SectionNameOrErr)
-      return SectionNameOrErr.takeError();
-    auto DecompressorOrErr =
-        Decompressor::create(*SectionNameOrErr, toStringRef(*ContentsOrErr),
-                             EF.isLE(), ELFT::Is64Bits);
-    if (!DecompressorOrErr)
-      return DecompressorOrErr.takeError();
-    size_t DecompressedSize = DecompressorOrErr->getDecompressedSize();
-    DecompressedContent = std::make_unique<uint8_t[]>(DecompressedSize);
-    MutableArrayRef<uint8_t> DecompressedContentRef(DecompressedContent.get(),
-                                                    DecompressedSize);
-    if (Error Err = DecompressorOrErr->decompress(DecompressedContentRef))
-      return std::move(Err);
-    Content = DecompressedContentRef;
+    // Read and optionally decompress section contents.
+    Expected<ArrayRef<uint8_t>> ContentsOrErr = EF.getSectionContents(Sec);
+    if (!ContentsOrErr)
+      return ContentsOrErr.takeError();
+    ArrayRef<uint8_t> Content = *ContentsOrErr;
+
+    std::unique_ptr<uint8_t[]> DecompressedContent;
+    if (Sec.sh_flags & llvm::ELF::SHF_COMPRESSED) {
+      Expected<StringRef> SectionNameOrErr = EF.getSectionName(Sec);
+      if (!SectionNameOrErr)
+        return SectionNameOrErr.takeError();
+      auto DecompressorOrErr =
+          Decompressor::create(*SectionNameOrErr, toStringRef(*ContentsOrErr),
+                               EF.isLE(), ELFT::Is64Bits);
+      if (!DecompressorOrErr)
+        return DecompressorOrErr.takeError();
+      size_t DecompressedSize = DecompressorOrErr->getDecompressedSize();
+      DecompressedContent = std::make_unique<uint8_t[]>(DecompressedSize);
+      MutableArrayRef<uint8_t> DecompressedContentRef(DecompressedContent.get(),
+                                                      DecompressedSize);
+      if (Error Err = DecompressorOrErr->decompress(DecompressedContentRef))
+        return std::move(Err);
+      Content = DecompressedContentRef;
+    }
+
+    DataExtractor Data(Content, EF.isLE(),
+                       sizeof(typename ELFFile<ELFT>::uintX_t));
+    return ELFBBAddrMapAddressExtractor(
+        std::move(Data), IsRelocatable, std::move(FunctionOffsetTranslations),
+        [&EF, &Sec] { return describe(EF, Sec); },
+        std::move(DecompressedContent));
   }
 
-  DataExtractor Data(Content, EF.isLE(),
-                     sizeof(typename ELFFile<ELFT>::uintX_t));
-  std::vector<BBAddrMap> FunctionEntries;
-
-  DataExtractor::Cursor Cur(0);
-  Error ULEBSizeErr = Error::success();
-  Error MetadataDecodeErr = Error::success();
+  std::string getSectionDescription() const override {
+    return DescribeSection ? DescribeSection() : "";
+  }
 
-  // Helper lambda to extract the (possibly relocatable) address stored at Cur.
-  auto ExtractAddress = [&]() -> Expected<uint64_t> {
-    uint64_t RelocationOffsetInSection = Cur.tell();
-    uint64_t Address = Data.getAddress(Cur);
-    if (!Cur)
-      return Cur.takeError();
-    if (!IsRelocatable)
-      return Address;
-    assert(Address == 0);
-    Expected<uint64_t> AddressOrErr =
-        GetAddressForRelocation(RelocationOffsetInSection);
+  Expected<uint64_t> extractAddress(DataExtractor::Cursor &Cur) override {
+    uint64_t Offset = Cur.tell();
+    Expected<uint64_t> AddressOrErr = AddressExtractor::extractAddress(Cur);
     if (!AddressOrErr)
       return AddressOrErr.takeError();
-    return *AddressOrErr;
-  };
-
-  uint8_t Version = 0;
-  uint16_t Feature = 0;
-  BBAddrMap::Features FeatEnable{};
-  while (!ULEBSizeErr && !MetadataDecodeErr && Cur &&
-         Cur.tell() < Content.size()) {
-    Version = Data.getU8(Cur);
-    if (!Cur)
-      break;
-    if (Version < 2 || Version > 5)
-      return createError("unsupported SHT_LLVM_BB_ADDR_MAP version: " +
-                         Twine(static_cast<int>(Version)));
-    Feature = Version < 5 ? Data.getU8(Cur) : Data.getU16(Cur);
-    if (!Cur)
-      break;
-    auto FeatEnableOrErr = BBAddrMap::Features::decode(Feature);
-    if (!FeatEnableOrErr)
-      return FeatEnableOrErr.takeError();
-    FeatEnable = *FeatEnableOrErr;
-    if (FeatEnable.CallsiteEndOffsets && Version < 3)
-      return createError("version should be >= 3 for SHT_LLVM_BB_ADDR_MAP when "
-                         "callsite offsets feature is enabled: version = " +
-                         Twine(static_cast<int>(Version)) +
-                         " feature = " + Twine(static_cast<int>(Feature)));
-    if (FeatEnable.BBHash && Version < 4)
-      return createError("version should be >= 4 for SHT_LLVM_BB_ADDR_MAP when "
-                         "basic block hash feature is enabled: version = " +
-                         Twine(static_cast<int>(Version)) +
-                         " feature = " + Twine(static_cast<int>(Feature)));
-    if (FeatEnable.PostLinkCfg && Version < 5)
-      return createError("version should be >= 5 for SHT_LLVM_BB_ADDR_MAP when "
-                         "post link cfg feature is enabled: version = " +
-                         Twine(static_cast<int>(Version)) +
-                         " feature = " + Twine(static_cast<int>(Feature)));
-    uint32_t NumBlocksInBBRange = 0;
-    uint32_t NumBBRanges = 1;
-    uint64_t RangeBaseAddress = 0;
-    if (FeatEnable.MultiBBRange) {
-      NumBBRanges = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-      if (!Cur || ULEBSizeErr)
-        break;
-      if (!NumBBRanges)
-        return createError("invalid zero number of BB ranges at offset " +
-                           Twine::utohexstr(Cur.tell()) + " in " +
-                           describe(EF, Sec));
-    } else {
-      auto AddressOrErr = ExtractAddress();
-      if (!AddressOrErr)
-        return AddressOrErr.takeError();
-      RangeBaseAddress = *AddressOrErr;
-      NumBlocksInBBRange = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-    }
-    std::vector<BBAddrMap::BBRangeEntry> BBRangeEntries;
-    uint32_t TotalNumBlocks = 0;
-    for (uint32_t BBRangeIndex = 0; BBRangeIndex < NumBBRanges;
-         ++BBRangeIndex) {
-      uint32_t PrevBBEndOffset = 0;
-      if (FeatEnable.MultiBBRange) {
-        auto AddressOrErr = ExtractAddress();
-        if (!AddressOrErr)
-          return AddressOrErr.takeError();
-        RangeBaseAddress = *AddressOrErr;
-        NumBlocksInBBRange = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-      }
-      std::vector<BBAddrMap::BBEntry> BBEntries;
-      if (!FeatEnable.OmitBBEntries) {
-        for (uint32_t BlockIndex = 0; !MetadataDecodeErr && !ULEBSizeErr &&
-                                      Cur && (BlockIndex < NumBlocksInBBRange);
-             ++BlockIndex) {
-          uint32_t ID = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-          uint32_t Offset = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-          // Read the callsite offsets.
-          uint32_t LastCallsiteEndOffset = 0;
-          SmallVector<uint32_t, 1> CallsiteEndOffsets;
-          if (FeatEnable.CallsiteEndOffsets) {
-            uint32_t NumCallsites =
-                readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-            CallsiteEndOffsets.reserve(NumCallsites);
-            for (uint32_t CallsiteIndex = 0;
-                 !ULEBSizeErr && Cur && (CallsiteIndex < NumCallsites);
-                 ++CallsiteIndex) {
-              LastCallsiteEndOffset +=
-                  readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-              CallsiteEndOffsets.push_back(LastCallsiteEndOffset);
-            }
-          }
-          uint32_t Size = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr) +
-                          LastCallsiteEndOffset;
-          uint32_t MD = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-          uint64_t Hash = FeatEnable.BBHash ? Data.getU64(Cur) : 0;
-          Expected<BBAddrMap::BBEntry::Metadata> MetadataOrErr =
-              BBAddrMap::BBEntry::Metadata::decode(MD);
-          if (!MetadataOrErr) {
-            MetadataDecodeErr = MetadataOrErr.takeError();
-            break;
-          }
-          BBEntries.push_back({ID, Offset + PrevBBEndOffset, Size,
-                               *MetadataOrErr, CallsiteEndOffsets, Hash});
-          PrevBBEndOffset += Offset + Size;
-        }
-        TotalNumBlocks += BBEntries.size();
-      }
-      BBRangeEntries.push_back({RangeBaseAddress, std::move(BBEntries)});
-    }
-    FunctionEntries.push_back({std::move(BBRangeEntries)});
-
-    if (PGOAnalyses || FeatEnable.hasPGOAnalysis()) {
-      // Function entry count
-      uint64_t FuncEntryCount =
-          FeatEnable.FuncEntryCount
-              ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
-              : 0;
-
-      std::vector<PGOAnalysisMap::PGOBBEntry> PGOBBEntries;
-      for (uint32_t BlockIndex = 0;
-           FeatEnable.hasPGOAnalysisBBData() && !MetadataDecodeErr &&
-           !ULEBSizeErr && Cur && (BlockIndex < TotalNumBlocks);
-           ++BlockIndex) {
-        // Block frequency
-        uint64_t BBF = FeatEnable.BBFreq
-                           ? readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr)
-                           : 0;
-        uint32_t PostLinkBBFreq =
-            FeatEnable.PostLinkCfg
-                ? readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr)
-                : 0;
-
-        // Branch probability
-        llvm::SmallVector<PGOAnalysisMap::PGOBBEntry::SuccessorEntry, 2>
-            Successors;
-        if (FeatEnable.BrProb) {
-          auto SuccCount = readULEB128As<uint64_t>(Data, Cur, ULEBSizeErr);
-          for (uint64_t I = 0; I < SuccCount; ++I) {
-            uint32_t BBID = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-            uint32_t BrProb = readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr);
-            uint32_t PostLinkFreq =
-                FeatEnable.PostLinkCfg
-                    ? readULEB128As<uint32_t>(Data, Cur, ULEBSizeErr)
-                    : 0;
-
-            if (PGOAnalyses)
-              Successors.push_back(
-                  {BBID, BranchProbability::getRaw(BrProb), PostLinkFreq});
-          }
-        }
-
-        if (PGOAnalyses)
-          PGOBBEntries.push_back(
-              {BlockFrequency(BBF), PostLinkBBFreq, std::move(Successors)});
-      }
-
-      if (PGOAnalyses)
-        PGOAnalyses->push_back(
-            {FuncEntryCount, std::move(PGOBBEntries), FeatEnable});
-    }
+    if (!IsRelocatable)
+      return *AddressOrErr;
+    assert(*AddressOrErr == 0);
+    auto FOTIterator = FunctionOffsetTranslations.find(Offset);
+    if (FOTIterator == FunctionOffsetTranslations.end())
+      return createError("failed to get relocation data for offset: " +
+                         Twine::utohexstr(Offset) + " in section " +
+                         getSectionDescription());
+    return FOTIterator->second;
   }
-  // Either Cur is in the error state, or we have an error in ULEBSizeErr or
-  // MetadataDecodeErr (but not both), but we join all errors here to be safe.
-  if (!Cur || ULEBSizeErr || MetadataDecodeErr)
-    return joinErrors(joinErrors(Cur.takeError(), std::move(ULEBSizeErr)),
-                      std::move(MetadataDecodeErr));
-  return FunctionEntries;
-}
+};
 
 template <class ELFT>
 Expected<std::vector<BBAddrMap>>
 ELFFile<ELFT>::decodeBBAddrMap(const Elf_Shdr &Sec, const Elf_Shdr *RelaSec,
                                std::vector<PGOAnalysisMap> *PGOAnalyses) const {
+  auto ExtractorOrErr =
+      ELFBBAddrMapAddressExtractor::create(*this, Sec, RelaSec);
+  if (!ExtractorOrErr)
+    return ExtractorOrErr.takeError();
+
   size_t OriginalPGOSize = PGOAnalyses ? PGOAnalyses->size() : 0;
-  auto AddrMapsOrErr = decodeBBAddrMapImpl(*this, Sec, RelaSec, PGOAnalyses);
+  auto AddrMapsOrErr =
+      llvm::object::decodeBBAddrMapPayload(*ExtractorOrErr, PGOAnalyses);
----------------
HaohaiWen wrote:

Done.

https://github.com/llvm/llvm-project/pull/188435


More information about the llvm-commits mailing list