[llvm-branch-commits] [llvm] release/20.x: [LoongArch] Fix assertion failure for annotate tablejump (#140907) (PR #141193)

[via llvm-branch-commits]

https://github.com/llvmbot created https://github.com/llvm/llvm-project/pull/141193

Backport 4e186f20e2f2be2fbf95d9713341a0b6507e707d

Requested by: @heiher

>From bc2bfeef77ad84512cec890f65944e46298dbd6c Mon Sep 17 00:00:00 2001
From: hev <wangrui at loongson.cn>
Date: Thu, 22 May 2025 18:50:40 +0800
Subject: [PATCH] [LoongArch] Fix assertion failure for annotate tablejump
 (#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 #140904

(cherry picked from commit 4e186f20e2f2be2fbf95d9713341a0b6507e707d)
---
 llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp  | 14 +++++++++-----
 .../LoongArch/LoongArchExpandPseudoInsts.cpp       |  3 ++-
 .../LoongArch/LoongArchMachineFunctionInfo.h       | 10 ++++------
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp b/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp
index 895a8e2646692..9a383f0a79a5c 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 c2d73a260b1c1..2107908be34ca 100644
--- a/llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp
@@ -638,7 +638,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