[BasicBlockSections] Using MBBSectionID as DenseMap key (llvm#97295)

getSectionIDNum may return same value for two different MBBSectionID.
e.g. A Cold type MBBSectionID with number 0 and a Default type
MBBSectionID with number 2 get same value 2 from getSectionIDNum. This
may lead to overwrite of MBBSectionRanges.  Using MBBSectionID itself
as DenseMap key is better choice.

Author: HaohaiWen

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -137,7 +137,7 @@ class AsmPrinter : public MachineFunctionPass {
     MCSymbol *BeginLabel, *EndLabel;
   };
 
-  MapVector<unsigned, MBBSectionRange> MBBSectionRanges;
+  MapVector<MBBSectionID, MBBSectionRange> MBBSectionRanges;
 
   /// Map global GOT equivalent MCSymbols to GlobalVariables and keep track of
   /// its number of uses by other globals.
@@ -157,7 +157,7 @@ class AsmPrinter : public MachineFunctionPass {
   /// Map a basic block section ID to the exception symbol associated with that
   /// section. Map entries are assigned and looked up via
   /// AsmPrinter::getMBBExceptionSym.
-  DenseMap<unsigned, MCSymbol *> MBBSectionExceptionSyms;
+  DenseMap<MBBSectionID, MCSymbol *> MBBSectionExceptionSyms;
 
   // The symbol used to represent the start of the current BB section of the
   // function. This is used to calculate the size of the BB section.

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_CODEGEN_MACHINEBASICBLOCK_H
 #define LLVM_CODEGEN_MACHINEBASICBLOCK_H
 
+#include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/SparseBitVector.h"
 #include "llvm/ADT/ilist.h"
@@ -74,6 +75,25 @@ struct MBBSectionID {
   MBBSectionID(SectionType T) : Type(T), Number(0) {}
 };
 
+template <> struct DenseMapInfo<MBBSectionID> {
+  using TypeInfo = DenseMapInfo<MBBSectionID::SectionType>;
+  using NumberInfo = DenseMapInfo<unsigned>;
+
+  static inline MBBSectionID getEmptyKey() {
+    return MBBSectionID(NumberInfo::getEmptyKey());
+  }
+  static inline MBBSectionID getTombstoneKey() {
+    return MBBSectionID(NumberInfo::getTombstoneKey());
+  }
+  static unsigned getHashValue(const MBBSectionID &SecID) {
+    return detail::combineHashValue(TypeInfo::getHashValue(SecID.Type),
+                                    NumberInfo::getHashValue(SecID.Number));
+  }
+  static bool isEqual(const MBBSectionID &LHS, const MBBSectionID &RHS) {
+    return LHS == RHS;
+  }
+};
+
 // This structure represents the information for a basic block pertaining to
 // the basic block sections profile.
 struct UniqueBBID {
@@ -658,12 +678,6 @@ class MachineBasicBlock
   /// Returns the section ID of this basic block.
   MBBSectionID getSectionID() const { return SectionID; }
 
-  /// Returns the unique section ID number of this basic block.
-  unsigned getSectionIDNum() const {
-    return ((unsigned)MBBSectionID::SectionType::Cold) -
-           ((unsigned)SectionID.Type) + SectionID.Number;
-  }
-
   /// Sets the fixed BBID of this basic block.
   void setBBID(const UniqueBBID &V) {
     assert(!BBID.has_value() && "Cannot change BBID.");

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1374,7 +1374,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
     OutStreamer->emitULEB128IntValue(MBBSectionRanges.size());
   }
   // Number of blocks in each MBB section.
-  MapVector<unsigned, unsigned> MBBSectionNumBlocks;
+  MapVector<MBBSectionID, unsigned> MBBSectionNumBlocks;
   const MCSymbol *PrevMBBEndSymbol = nullptr;
   if (!Features.MultiBBRange) {
     OutStreamer->AddComment("function address");
@@ -1388,7 +1388,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
       BBCount++;
       if (MBB.isEndSection()) {
         // Store each section's basic block count when it ends.
-        MBBSectionNumBlocks[MBB.getSectionIDNum()] = BBCount;
+        MBBSectionNumBlocks[MBB.getSectionID()] = BBCount;
         // Reset the count for the next section.
         BBCount = 0;
       }
@@ -1404,8 +1404,7 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
       OutStreamer->AddComment("base address");
       OutStreamer->emitSymbolValue(MBBSymbol, getPointerSize());
       OutStreamer->AddComment("number of basic blocks");
-      OutStreamer->emitULEB128IntValue(
-          MBBSectionNumBlocks[MBB.getSectionIDNum()]);
+      OutStreamer->emitULEB128IntValue(MBBSectionNumBlocks[MBB.getSectionID()]);
       PrevMBBEndSymbol = MBBSymbol;
     }
     // TODO: Remove this check when version 1 is deprecated.
@@ -1855,7 +1854,9 @@ void AsmPrinter::emitFunctionBody() {
               OutContext);
           OutStreamer->emitELFSize(CurrentSectionBeginSym, SizeExp);
         }
-        MBBSectionRanges[MBB.getSectionIDNum()] =
+        assert(!MBBSectionRanges.contains(MBB.getSectionID()) &&
+               "Overwrite section range");
+        MBBSectionRanges[MBB.getSectionID()] =
             MBBSectionRange{CurrentSectionBeginSym, MBB.getEndSymbol()};
       }
     }
@@ -1972,7 +1973,9 @@ void AsmPrinter::emitFunctionBody() {
   for (auto &Handler : Handlers)
     Handler->markFunctionEnd();
 
-  MBBSectionRanges[MF->front().getSectionIDNum()] =
+  assert(!MBBSectionRanges.contains(MF->front().getSectionID()) &&
+         "Overwrite section range");
+  MBBSectionRanges[MF->front().getSectionID()] =
       MBBSectionRange{CurrentFnBegin, CurrentFnEnd};
 
   // Print out jump tables referenced by the function.
@@ -2536,7 +2539,7 @@ bool AsmPrinter::doFinalization(Module &M) {
 }
 
 MCSymbol *AsmPrinter::getMBBExceptionSym(const MachineBasicBlock &MBB) {
-  auto Res = MBBSectionExceptionSyms.try_emplace(MBB.getSectionIDNum());
+  auto Res = MBBSectionExceptionSyms.try_emplace(MBB.getSectionID());
   if (Res.second)
     Res.first->second = createTempSymbol("exception");
   return Res.first->second;

llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.cpp

@@ -684,7 +684,7 @@ void DwarfCompileUnit::attachRangesOrLowHighPC(
     // the order of blocks will be frozen beyond this point.
     do {
       if (MBB->sameSection(EndMBB) || MBB->isEndSection()) {
-        auto MBBSectionRange = Asm->MBBSectionRanges[MBB->getSectionIDNum()];
+        auto MBBSectionRange = Asm->MBBSectionRanges[MBB->getSectionID()];
         List.push_back(
             {MBB->sameSection(BeginMBB) ? BeginLabel
                                         : MBBSectionRange.BeginLabel,

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp

@@ -1713,7 +1713,7 @@ bool DwarfDebug::buildLocationList(SmallVectorImpl<DebugLocEntry> &DebugLoc,
     const MCSymbol *EndLabel;
     if (std::next(EI) == Entries.end()) {
       const MachineBasicBlock &EndMBB = Asm->MF->back();
-      EndLabel = Asm->MBBSectionRanges[EndMBB.getSectionIDNum()].EndLabel;
+      EndLabel = Asm->MBBSectionRanges[EndMBB.getSectionID()].EndLabel;
       if (EI->isClobber())
         EndMI = EI->getInstr();
     }
@@ -2064,7 +2064,7 @@ void DwarfDebug::beginInstruction(const MachineInstr *MI) {
 
   bool PrevInstInSameSection =
       (!PrevInstBB ||
-       PrevInstBB->getSectionIDNum() == MI->getParent()->getSectionIDNum());
+       PrevInstBB->getSectionID() == MI->getParent()->getSectionID());
   if (DL == PrevInstLoc && PrevInstInSameSection) {
     // If we have an ongoing unspecified location, nothing to do here.
     if (!DL)

llvm/lib/CodeGen/AsmPrinter/EHStreamer.cpp

@@ -253,8 +253,8 @@ void EHStreamer::computeCallSiteTable(
       // We start a call-site range upon function entry and at the beginning of
       // every basic block section.
       CallSiteRanges.push_back(
-          {Asm->MBBSectionRanges[MBB.getSectionIDNum()].BeginLabel,
-           Asm->MBBSectionRanges[MBB.getSectionIDNum()].EndLabel,
+          {Asm->MBBSectionRanges[MBB.getSectionID()].BeginLabel,
+           Asm->MBBSectionRanges[MBB.getSectionID()].EndLabel,
            Asm->getMBBExceptionSym(MBB), CallSites.size()});
       PreviousIsInvoke = false;
       SawPotentiallyThrowing = false;