[LoongArch] Fix assertion failure for annotate tablejump

[heiher]

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 #140904

[llvmbot]

@llvm/pr-subscribers-backend-loongarch

Author: hev (heiher)

Changes

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 #140904

---

Full diff: https://github.com/llvm/llvm-project/pull/140907.diff

3 Files Affected:

• (modified) llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp (+9-5)
• (modified) llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp (+2-1)
• (modified) llvm/lib/Target/LoongArch/LoongArchMachineFunctionInfo.h (+4-6)

diff --git a/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp b/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp
index 0672570c4dcd0..9181e539f75cb 100644
--- a/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp
@@ -265,13 +265,16 @@ 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.
@@ -279,14 +282,15 @@ void LoongArchAsmPrinter::emitJumpTableInfo() {
       ".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);
   }
 }
 
diff --git a/llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp b/llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp
index 9e9c65a041bf7..7aef4ab53e4ea 100644
--- a/llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp
+++ b/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;
         }
diff --git a/llvm/lib/Target/LoongArch/LoongArchMachineFunctionInfo.h b/llvm/lib/Target/LoongArch/LoongArchMachineFunctionInfo.h
index daa47c4dc7e32..904985c189dba 100644
--- a/llvm/lib/Target/LoongArch/LoongArchMachineFunctionInfo.h
+++ b/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

[wangleiat]

LGTM, thanks.

[heiher]

/cherry-pick 4e186f2

[llvmbot]

/pull-request #141193