[LoongArch] Fix assertion failure for annotate tablejump (llvm#140907)

Fix a use-after-free issue related to annotateTableJump in the LoongArch
target.

Previously, `LoongArchPreRAExpandPseudo::annotateTableJump()` recorded a
reference to a MachineOperand representing a jump table index. However,
later optimizations such as the `BranchFolder` pass may delete the
instruction containing this operand, leaving a dangling reference.

This led to an assertion failure in
`LoongArchAsmPrinter::emitJumpTableInfo()` when trying to access a freed
MachineOperand via `getIndex()`.

The fix avoids holding a reference to the MachineOperand. Instead, we
extract and store the jump table index at the time of annotation. During
`emitJumpTableInfo()`, we verify whether the recorded index still exists
in the MachineFunction's jump table. If not, we skip emission for that
entry.

Fixes llvm#140904

Author: heiher

llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp

@@ -265,28 +265,32 @@ void LoongArchAsmPrinter::emitJumpTableInfo() {
 
   assert(TM.getTargetTriple().isOSBinFormatELF());
 
-  unsigned Size = getDataLayout().getPointerSize();
   auto *LAFI = MF->getInfo<LoongArchMachineFunctionInfo>();
   unsigned EntrySize = LAFI->getJumpInfoSize();
+  auto JTI = MF->getJumpTableInfo();
 
-  if (0 == EntrySize)
+  if (!JTI || 0 == EntrySize)
     return;
 
+  unsigned Size = getDataLayout().getPointerSize();
+  auto JT = JTI->getJumpTables();
+
   // Emit an additional section to store the correlation info as pairs of
   // addresses, each pair contains the address of a jump instruction (jr) and
   // the address of the jump table.
   OutStreamer->switchSection(MMI->getContext().getELFSection(
       ".discard.tablejump_annotate", ELF::SHT_PROGBITS, 0));
 
   for (unsigned Idx = 0; Idx < EntrySize; ++Idx) {
+    int JTIIdx = LAFI->getJumpInfoJTIIndex(Idx);
+    if (JT[JTIIdx].MBBs.empty())
+      continue;
     OutStreamer->emitValue(
         MCSymbolRefExpr::create(LAFI->getJumpInfoJrMI(Idx)->getPreInstrSymbol(),
                                 OutContext),
         Size);
     OutStreamer->emitValue(
-        MCSymbolRefExpr::create(
-            GetJTISymbol(LAFI->getJumpInfoJTIMO(Idx)->getIndex()), OutContext),
-        Size);
+        MCSymbolRefExpr::create(GetJTISymbol(JTIIdx), OutContext), Size);
   }
 }
 

llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp

@@ -636,7 +636,8 @@ void LoongArchPreRAExpandPseudo::annotateTableJump(
         if (MO.isJTI()) {
           MBBI->setPreInstrSymbol(
               *MF, MF->getContext().createNamedTempSymbol("jrtb_"));
-          MF->getInfo<LoongArchMachineFunctionInfo>()->setJumpInfo(&*MBBI, &MO);
+          MF->getInfo<LoongArchMachineFunctionInfo>()->setJumpInfo(
+              &*MBBI, MO.getIndex());
           IsFound = true;
           return;
         }

llvm/lib/Target/LoongArch/LoongArchMachineFunctionInfo.h

@@ -41,7 +41,7 @@ class LoongArchMachineFunctionInfo : public MachineFunctionInfo {
 
   /// Pairs of `jr` instructions and corresponding JTI operands, used for the
   /// `annotate-tablejump` option.
-  SmallVector<std::pair<MachineInstr *, MachineOperand *>, 4> JumpInfos;
+  SmallVector<std::pair<MachineInstr *, int>, 4> JumpInfos;
 
 public:
   LoongArchMachineFunctionInfo(const Function &F,
@@ -76,14 +76,12 @@ class LoongArchMachineFunctionInfo : public MachineFunctionInfo {
     return is_contained(SExt32Registers, Reg);
   }
 
-  void setJumpInfo(MachineInstr *JrMI, MachineOperand *JTIMO) {
-    JumpInfos.push_back(std::make_pair(JrMI, JTIMO));
+  void setJumpInfo(MachineInstr *JrMI, int JTIIdx) {
+    JumpInfos.push_back(std::make_pair(JrMI, JTIIdx));
   }
   unsigned getJumpInfoSize() { return JumpInfos.size(); }
   MachineInstr *getJumpInfoJrMI(unsigned Idx) { return JumpInfos[Idx].first; }
-  MachineOperand *getJumpInfoJTIMO(unsigned Idx) {
-    return JumpInfos[Idx].second;
-  }
+  int getJumpInfoJTIIndex(unsigned Idx) { return JumpInfos[Idx].second; }
 };
 
 } // end namespace llvm