1 files changed, 132 insertions, 74 deletions

diff --git a/llvm/tools/llvm-objdump/llvm-objdump.cpp b/llvm/tools/llvm-objdump/llvm-objdump.cpp
index 7f57713..b4467ec 100644
--- a/llvm/tools/llvm-objdump/llvm-objdump.cpp
+++ b/llvm/tools/llvm-objdump/llvm-objdump.cpp
@@ -172,6 +172,110 @@ public:
                        "Mach-O object file displaying tool") {}
 };
 
+struct BBAddrMapLabel {
+  std::string BlockLabel;
+  std::string PGOAnalysis;
+};
+
+// This class represents the BBAddrMap and PGOMap associated with a single
+// function.
+class BBAddrMapFunctionEntry {
+public:
+  BBAddrMapFunctionEntry(BBAddrMap AddrMap, PGOAnalysisMap PGOMap)
+      : AddrMap(std::move(AddrMap)), PGOMap(std::move(PGOMap)) {}
+
+  const BBAddrMap &getAddrMap() const { return AddrMap; }
+
+  // Returns the PGO string associated with the entry of index `PGOBBEntryIndex`
+  // in `PGOMap`.
+  std::string constructPGOLabelString(size_t PGOBBEntryIndex) const {
+    if (!PGOMap.FeatEnable.hasPGOAnalysis())
+      return "";
+    std::string PGOString;
+    raw_string_ostream PGOSS(PGOString);
+
+    PGOSS << " (";
+    if (PGOMap.FeatEnable.FuncEntryCount && PGOBBEntryIndex == 0) {
+      PGOSS << "Entry count: " << Twine(PGOMap.FuncEntryCount);
+      if (PGOMap.FeatEnable.hasPGOAnalysisBBData()) {
+        PGOSS << ", ";
+      }
+    }
+
+    if (PGOMap.FeatEnable.hasPGOAnalysisBBData()) {
+
+      assert(PGOBBEntryIndex < PGOMap.BBEntries.size() &&
+             "Expected PGOAnalysisMap and BBAddrMap to have the same entries");
+      const PGOAnalysisMap::PGOBBEntry &PGOBBEntry =
+          PGOMap.BBEntries[PGOBBEntryIndex];
+
+      if (PGOMap.FeatEnable.BBFreq) {
+        PGOSS << "Frequency: " << Twine(PGOBBEntry.BlockFreq.getFrequency());
+        if (PGOMap.FeatEnable.BrProb && PGOBBEntry.Successors.size() > 0) {
+          PGOSS << ", ";
+        }
+      }
+      if (PGOMap.FeatEnable.BrProb && PGOBBEntry.Successors.size() > 0) {
+        PGOSS << "Successors: ";
+        interleaveComma(
+            PGOBBEntry.Successors, PGOSS,
+            [&PGOSS](const PGOAnalysisMap::PGOBBEntry::SuccessorEntry &SE) {
+              PGOSS << "BB" << SE.ID << ":";
+              PGOSS.write_hex(SE.Prob.getNumerator());
+            });
+      }
+    }
+    PGOSS << ")";
+
+    return PGOString;
+  }
+
+private:
+  const BBAddrMap AddrMap;
+  const PGOAnalysisMap PGOMap;
+};
+
+// This class represents the BBAddrMap and PGOMap of potentially multiple
+// functions in a section.
+class BBAddrMapInfo {
+public:
+  void clear() {
+    FunctionAddrToMap.clear();
+    RangeBaseAddrToFunctionAddr.clear();
+  }
+
+  bool empty() const { return FunctionAddrToMap.empty(); }
+
+  void AddFunctionEntry(BBAddrMap AddrMap, PGOAnalysisMap PGOMap) {
+    uint64_t FunctionAddr = AddrMap.getFunctionAddress();
+    for (size_t I = 1; I < AddrMap.BBRanges.size(); ++I)
+      RangeBaseAddrToFunctionAddr.emplace(AddrMap.BBRanges[I].BaseAddress,
+                                          FunctionAddr);
+    [[maybe_unused]] auto R = FunctionAddrToMap.try_emplace(
+        FunctionAddr, std::move(AddrMap), std::move(PGOMap));
+    assert(R.second && "duplicate function address");
+  }
+
+  // Returns the BBAddrMap entry for the function associated with `BaseAddress`.
+  // `BaseAddress` could be the function address or the address of a range
+  // associated with that function. Returns `nullptr` if `BaseAddress` is not
+  // mapped to any entry.
+  const BBAddrMapFunctionEntry *getEntryForAddress(uint64_t BaseAddress) const {
+    uint64_t FunctionAddr = BaseAddress;
+    auto S = RangeBaseAddrToFunctionAddr.find(BaseAddress);
+    if (S != RangeBaseAddrToFunctionAddr.end())
+      FunctionAddr = S->second;
+    auto R = FunctionAddrToMap.find(FunctionAddr);
+    if (R == FunctionAddrToMap.end())
+      return nullptr;
+    return &R->second;
+  }
+
+private:
+  std::unordered_map<uint64_t, BBAddrMapFunctionEntry> FunctionAddrToMap;
+  std::unordered_map<uint64_t, uint64_t> RangeBaseAddrToFunctionAddr;
+};
+
 } // namespace
 
 #define DEBUG_TYPE "objdump"
@@ -1266,80 +1370,38 @@ static SymbolInfoTy createDummySymbolInfo(const ObjectFile &Obj,
     return SymbolInfoTy(Addr, Name, Type);
 }
 
-struct BBAddrMapLabel {
-  std::string BlockLabel;
-  std::string PGOAnalysis;
-};
-
-static std::string constructPGOLabelString(const PGOAnalysisMap &PGOMap,
-                                           size_t BBEntryIndex) {
-  std::string PGOString;
-  raw_string_ostream PGOSS(PGOString);
-
-  PGOSS << " (";
-  if (PGOMap.FeatEnable.FuncEntryCount && BBEntryIndex == 0) {
-    PGOSS << "Entry count: " << Twine(PGOMap.FuncEntryCount);
-    if (PGOMap.FeatEnable.BBFreq || PGOMap.FeatEnable.BrProb) {
-      PGOSS << ", ";
-    }
-  }
-
-  if (PGOMap.FeatEnable.BBFreq || PGOMap.FeatEnable.BrProb) {
-    assert(BBEntryIndex < PGOMap.BBEntries.size() &&
-           "Expected PGOAnalysisMap and BBAddrMap to have the same entires");
-    const PGOAnalysisMap::PGOBBEntry &PGOBBEntry =
-        PGOMap.BBEntries[BBEntryIndex];
-
-    if (PGOMap.FeatEnable.BBFreq) {
-      PGOSS << "Frequency: " << Twine(PGOBBEntry.BlockFreq.getFrequency());
-      if (PGOMap.FeatEnable.BrProb && PGOBBEntry.Successors.size() > 0) {
-        PGOSS << ", ";
-      }
-    }
-    if (PGOMap.FeatEnable.BrProb && PGOBBEntry.Successors.size() > 0) {
-      PGOSS << "Successors: ";
-      interleaveComma(
-          PGOBBEntry.Successors, PGOSS,
-          [&PGOSS](const PGOAnalysisMap::PGOBBEntry::SuccessorEntry &SE) {
-            PGOSS << "BB" << SE.ID << ":";
-            PGOSS.write_hex(SE.Prob.getNumerator());
-          });
-    }
-  }
-  PGOSS << ")";
-
-  return PGOString;
-}
-
 static void collectBBAddrMapLabels(
-    const std::unordered_map<uint64_t, BBAddrMap> &AddrToBBAddrMap,
-    const std::unordered_map<uint64_t, PGOAnalysisMap> &AddrToPGOAnalysisMap,
-    uint64_t SectionAddr, uint64_t Start, uint64_t End,
-    std::unordered_map<uint64_t, std::vector<BBAddrMapLabel>> &Labels,
-    const StringRef FileName) {
-  if (AddrToBBAddrMap.empty())
+    const BBAddrMapInfo &FullAddrMap, uint64_t SectionAddr, uint64_t Start,
+    uint64_t End,
+    std::unordered_map<uint64_t, std::vector<BBAddrMapLabel>> &Labels) {
+  if (FullAddrMap.empty())
     return;
   Labels.clear();
   uint64_t StartAddress = SectionAddr + Start;
   uint64_t EndAddress = SectionAddr + End;
-  auto Iter = AddrToBBAddrMap.find(StartAddress);
-  if (Iter == AddrToBBAddrMap.end())
+  const BBAddrMapFunctionEntry *FunctionMap =
+      FullAddrMap.getEntryForAddress(StartAddress);
+  if (!FunctionMap)
     return;
-  auto PGOIter = AddrToPGOAnalysisMap.find(StartAddress);
-
-  for (size_t I = 0; I < Iter->second.getBBEntries().size(); ++I) {
-    const BBAddrMap::BBEntry &BBEntry = Iter->second.getBBEntries()[I];
-    uint64_t BBAddress = BBEntry.Offset + Iter->second.getFunctionAddress();
+  std::optional<size_t> BBRangeIndex =
+      FunctionMap->getAddrMap().getBBRangeIndexForBaseAddress(StartAddress);
+  if (!BBRangeIndex)
+    return;
+  size_t NumBBEntriesBeforeRange = 0;
+  for (size_t I = 0; I < *BBRangeIndex; ++I)
+    NumBBEntriesBeforeRange +=
+        FunctionMap->getAddrMap().BBRanges[I].BBEntries.size();
+  const auto &BBRange = FunctionMap->getAddrMap().BBRanges[*BBRangeIndex];
+  for (size_t I = 0; I < BBRange.BBEntries.size(); ++I) {
+    const BBAddrMap::BBEntry &BBEntry = BBRange.BBEntries[I];
+    uint64_t BBAddress = BBEntry.Offset + BBRange.BaseAddress;
     if (BBAddress >= EndAddress)
       continue;
 
     std::string LabelString = ("BB" + Twine(BBEntry.ID)).str();
-    std::string PGOString;
-
-    if (PGOIter != AddrToPGOAnalysisMap.end())
-      PGOString = constructPGOLabelString(PGOIter->second, I);
-
-    Labels[BBAddress].push_back({LabelString, PGOString});
+    Labels[BBAddress].push_back(
+        {LabelString,
+         FunctionMap->constructPGOLabelString(NumBBEntriesBeforeRange + I)});
   }
 }
 
@@ -1696,11 +1758,10 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
 
   LLVM_DEBUG(LVP.dump());
 
-  std::unordered_map<uint64_t, BBAddrMap> AddrToBBAddrMap;
-  std::unordered_map<uint64_t, PGOAnalysisMap> AddrToPGOAnalysisMap;
+  BBAddrMapInfo FullAddrMap;
   auto ReadBBAddrMap = [&](std::optional<unsigned> SectionIndex =
                                std::nullopt) {
-    AddrToBBAddrMap.clear();
+    FullAddrMap.clear();
     if (const auto *Elf = dyn_cast<ELFObjectFileBase>(&Obj)) {
       std::vector<PGOAnalysisMap> PGOAnalyses;
       auto BBAddrMapsOrErr = Elf->readBBAddrMap(SectionIndex, &PGOAnalyses);
@@ -1710,10 +1771,8 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
       }
       for (auto &&[FunctionBBAddrMap, FunctionPGOAnalysis] :
            zip_equal(*std::move(BBAddrMapsOrErr), std::move(PGOAnalyses))) {
-        uint64_t Addr = FunctionBBAddrMap.Addr;
-        AddrToBBAddrMap.emplace(Addr, std::move(FunctionBBAddrMap));
-        if (FunctionPGOAnalysis.FeatEnable.anyEnabled())
-          AddrToPGOAnalysisMap.emplace(Addr, std::move(FunctionPGOAnalysis));
+        FullAddrMap.AddFunctionEntry(std::move(FunctionBBAddrMap),
+                                     std::move(FunctionPGOAnalysis));
       }
     }
   };
@@ -2049,9 +2108,8 @@ disassembleObject(ObjectFile &Obj, const ObjectFile &DbgObj,
                                   DT->DisAsm.get(), DT->InstPrinter.get(),
                                   PrimaryTarget.SubtargetInfo.get(),
                                   SectionAddr, Index, End, AllLabels);
-        collectBBAddrMapLabels(AddrToBBAddrMap, AddrToPGOAnalysisMap,
-                               SectionAddr, Index, End, BBAddrMapLabels,
-                               FileName);
+        collectBBAddrMapLabels(FullAddrMap, SectionAddr, Index, End,
+                               BBAddrMapLabels);
       }
 
       if (DT->InstrAnalysis)