[CodeGen][DebugInfo] Add missing debug info for jump table BB (#71021)

visitJumpTable is called on FinishBasicBlock. At that time, getCurSDLoc
will always return SDLoc without DebugLoc since CurInst was set to
nullptr after visiting each instruction.
This patch passes SDLoc to buildJumpTable when visiting SwitchInst so
that visitJumpTable can use it later.

5 files changed, 24 insertions, 17 deletions

diff --git a/llvm/include/llvm/CodeGen/SwitchLoweringUtils.h b/llvm/include/llvm/CodeGen/SwitchLoweringUtils.h
index 189dfef590b0..5d06e21737b8 100644
--- a/llvm/include/llvm/CodeGen/SwitchLoweringUtils.h
+++ b/llvm/include/llvm/CodeGen/SwitchLoweringUtils.h
@@ -174,8 +174,12 @@ struct JumpTable {
   /// check MBB.  This is when updating PHI nodes in successors.
   MachineBasicBlock *Default;
 
-  JumpTable(unsigned R, unsigned J, MachineBasicBlock *M, MachineBasicBlock *D)
-      : Reg(R), JTI(J), MBB(M), Default(D) {}
+  /// The debug location of the instruction this JumpTable was produced from.
+  std::optional<SDLoc> SL; // For SelectionDAG
+
+  JumpTable(unsigned R, unsigned J, MachineBasicBlock *M, MachineBasicBlock *D,
+            std::optional<SDLoc> SL)
+      : Reg(R), JTI(J), MBB(M), Default(D), SL(SL) {}
 };
 struct JumpTableHeader {
   APInt First;
@@ -270,14 +274,14 @@ public:
   std::vector<BitTestBlock> BitTestCases;
 
   void findJumpTables(CaseClusterVector &Clusters, const SwitchInst *SI,
-                      MachineBasicBlock *DefaultMBB,
+                      std::optional<SDLoc> SL, MachineBasicBlock *DefaultMBB,
                       ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI);
 
   bool buildJumpTable(const CaseClusterVector &Clusters, unsigned First,
                       unsigned Last, const SwitchInst *SI,
+                      const std::optional<SDLoc> &SL,
                       MachineBasicBlock *DefaultMBB, CaseCluster &JTCluster);
 
-
   void findBitTestClusters(CaseClusterVector &Clusters, const SwitchInst *SI);
 
   /// Build a bit test cluster from Clusters[First..Last]. Returns false if it
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index d7ea25838058..62450e4c43ff 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -722,7 +722,7 @@ bool IRTranslator::translateSwitch(const User &U, MachineIRBuilder &MIB) {
     return true;
   }
 
-  SL->findJumpTables(Clusters, &SI, DefaultMBB, nullptr, nullptr);
+  SL->findJumpTables(Clusters, &SI, std::nullopt, DefaultMBB, nullptr, nullptr);
   SL->findBitTestClusters(Clusters, &SI);
 
   LLVM_DEBUG({
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index f0021f9f0907..2fa37f4c519e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -2670,14 +2670,13 @@ void SelectionDAGBuilder::visitSwitchCase(CaseBlock &CB,
 /// visitJumpTable - Emit JumpTable node in the current MBB
 void SelectionDAGBuilder::visitJumpTable(SwitchCG::JumpTable &JT) {
   // Emit the code for the jump table
+  assert(JT.SL && "Should set SDLoc for SelectionDAG!");
   assert(JT.Reg != -1U && "Should lower JT Header first!");
   EVT PTy = DAG.getTargetLoweringInfo().getPointerTy(DAG.getDataLayout());
-  SDValue Index = DAG.getCopyFromReg(getControlRoot(), getCurSDLoc(),
-                                     JT.Reg, PTy);
+  SDValue Index = DAG.getCopyFromReg(getControlRoot(), *JT.SL, JT.Reg, PTy);
   SDValue Table = DAG.getJumpTable(JT.JTI, PTy);
-  SDValue BrJumpTable = DAG.getNode(ISD::BR_JT, getCurSDLoc(),
-                                    MVT::Other, Index.getValue(1),
-                                    Table, Index);
+  SDValue BrJumpTable = DAG.getNode(ISD::BR_JT, *JT.SL, MVT::Other,
+                                    Index.getValue(1), Table, Index);
   DAG.setRoot(BrJumpTable);
 }
 
@@ -2686,7 +2685,8 @@ void SelectionDAGBuilder::visitJumpTable(SwitchCG::JumpTable &JT) {
 void SelectionDAGBuilder::visitJumpTableHeader(SwitchCG::JumpTable &JT,
                                                JumpTableHeader &JTH,
                                                MachineBasicBlock *SwitchBB) {
-  SDLoc dl = getCurSDLoc();
+  assert(JT.SL && "Should set SDLoc for SelectionDAG!");
+  const SDLoc &dl = *JT.SL;
 
   // Subtract the lowest switch case value from the value being switched on.
   SDValue SwitchOp = getValue(JTH.SValue);
@@ -11869,7 +11869,8 @@ void SelectionDAGBuilder::visitSwitch(const SwitchInst &SI) {
     return;
   }
 
-  SL->findJumpTables(Clusters, &SI, DefaultMBB, DAG.getPSI(), DAG.getBFI());
+  SL->findJumpTables(Clusters, &SI, getCurSDLoc(), DefaultMBB, DAG.getPSI(),
+                     DAG.getBFI());
   SL->findBitTestClusters(Clusters, &SI);
 
   LLVM_DEBUG({
diff --git a/llvm/lib/CodeGen/SwitchLoweringUtils.cpp b/llvm/lib/CodeGen/SwitchLoweringUtils.cpp
index b01a8bed0a39..7982d80353bd 100644
--- a/llvm/lib/CodeGen/SwitchLoweringUtils.cpp
+++ b/llvm/lib/CodeGen/SwitchLoweringUtils.cpp
@@ -45,6 +45,7 @@ SwitchCG::getJumpTableNumCases(const SmallVectorImpl<unsigned> &TotalCases,
 
 void SwitchCG::SwitchLowering::findJumpTables(CaseClusterVector &Clusters,
                                               const SwitchInst *SI,
+                                              std::optional<SDLoc> SL,
                                               MachineBasicBlock *DefaultMBB,
                                               ProfileSummaryInfo *PSI,
                                               BlockFrequencyInfo *BFI) {
@@ -87,7 +88,7 @@ void SwitchCG::SwitchLowering::findJumpTables(CaseClusterVector &Clusters,
   // Cheap case: the whole range may be suitable for jump table.
   if (TLI->isSuitableForJumpTable(SI, NumCases, Range, PSI, BFI)) {
     CaseCluster JTCluster;
-    if (buildJumpTable(Clusters, 0, N - 1, SI, DefaultMBB, JTCluster)) {
+    if (buildJumpTable(Clusters, 0, N - 1, SI, SL, DefaultMBB, JTCluster)) {
       Clusters[0] = JTCluster;
       Clusters.resize(1);
       return;
@@ -177,7 +178,7 @@ void SwitchCG::SwitchLowering::findJumpTables(CaseClusterVector &Clusters,
 
     CaseCluster JTCluster;
     if (NumClusters >= MinJumpTableEntries &&
-        buildJumpTable(Clusters, First, Last, SI, DefaultMBB, JTCluster)) {
+        buildJumpTable(Clusters, First, Last, SI, SL, DefaultMBB, JTCluster)) {
       Clusters[DstIndex++] = JTCluster;
     } else {
       for (unsigned I = First; I <= Last; ++I)
@@ -190,6 +191,7 @@ void SwitchCG::SwitchLowering::findJumpTables(CaseClusterVector &Clusters,
 bool SwitchCG::SwitchLowering::buildJumpTable(const CaseClusterVector &Clusters,
                                               unsigned First, unsigned Last,
                                               const SwitchInst *SI,
+                                              const std::optional<SDLoc> &SL,
                                               MachineBasicBlock *DefaultMBB,
                                               CaseCluster &JTCluster) {
   assert(First <= Last);
@@ -251,7 +253,7 @@ bool SwitchCG::SwitchLowering::buildJumpTable(const CaseClusterVector &Clusters,
                      ->createJumpTableIndex(Table);
 
   // Set up the jump table info.
-  JumpTable JT(-1U, JTI, JumpTableMBB, nullptr);
+  JumpTable JT(-1U, JTI, JumpTableMBB, nullptr, SL);
   JumpTableHeader JTH(Clusters[First].Low->getValue(),
                       Clusters[Last].High->getValue(), SI->getCondition(),
                       nullptr, false);
diff --git a/llvm/test/DebugInfo/X86/debug-info-jump-table.ll b/llvm/test/DebugInfo/X86/debug-info-jump-table.ll
index f60ca32a0671..067b80bb5eab 100644
--- a/llvm/test/DebugInfo/X86/debug-info-jump-table.ll
+++ b/llvm/test/DebugInfo/X86/debug-info-jump-table.ll
@@ -24,8 +24,8 @@ define dso_local void @foo(i32 noundef %cond) local_unnamed_addr #0 !dbg !42 {
 ;CHECK: Initial selection DAG: %bb.{{[0-9]+}} 'foo:entry'
 ;CHECK: SelectionDAG has 5 nodes:
 ;CHECK:     [[TMP1:t.*]]: ch,glue = EntryToken
-;CHECK:   [[TMP2:t.*]]: i64,ch = CopyFromReg [[TMP1]], Register:i64 %{{[0-9]+}}
-;CHECK:   t{{[0-9]+}}: ch = br_jt [[TMP2]]:1, JumpTable:i64<0>, [[TMP2]]
+;CHECK:   [[TMP2:t.*]]: i64,ch = CopyFromReg [[TMP1]], Register:i64 %{{[0-9]+}}, jump_table.c:4:3
+;CHECK:   t{{[0-9]+}}: ch = br_jt [[TMP2]]:1, JumpTable:i64<0>, [[TMP2]], jump_table.c:4:3
 
 entry:
   call void @llvm.dbg.value(metadata i32 %cond, metadata !47, metadata !DIExpression()), !dbg !48