Fix the compilation assertion due to unreachable BB pruning not deleting the associated BB from the jump tables

[NeHuang]

Summary

We encountered the same failure as https://reviews.llvm.org/D22435. Original bug tracked in https://bugs.llvm.org/show_bug.cgi?id=28102
@cperfetto provided a potential solution to remove the unreachable MBBs reference in the jump table. (Comment 27 and Comment 28)

We have confirmed the assertion issue fixed with this patch.

[lebedev.ri]

Test missing

[NeHuang]

Add the test case.

[NeHuang]

ping

[bsaleil]

LGTM

[amyk]

Minor nit but LGTM

llvm/lib/CodeGen/MachineFunction.cpp

1053

nit: end with period (as well as the other comments in this patch that are added).

[NeHuang]

Addressed review comments for the nits.

llvm/include/llvm/CodeGen/MachineJumpTableInfo.h

public:
   }
 
   /// RemoveJumpTable - Mark the specific index as being dead.  This will
   /// prevent it from being emitted.
   void RemoveJumpTable(unsigned Idx) {
     JumpTables[Idx].MBBs.clear();
   }
 
+  /// RemoveMBBFromJumpTables - If MBB is present in any jump tables, remove it
+  bool RemoveMBBFromJumpTables(MachineBasicBlock *MBB);
+
   /// ReplaceMBBInJumpTables - If Old is the target of any jump tables, update
   /// the jump tables to branch to New instead.
   bool ReplaceMBBInJumpTables(MachineBasicBlock *Old, MachineBasicBlock *New);
 
   /// ReplaceMBBInJumpTable - If Old is a target of the jump tables, update
   /// the jump table to branch to New instead.
   bool ReplaceMBBInJumpTable(unsigned Idx, MachineBasicBlock *Old,
                              MachineBasicBlock *New);

llvm/lib/CodeGen/MachineFunction.cpp

MachineFunction::CreateMachineBasicBlock(const BasicBlock *bb) {
   return new (BasicBlockRecycler.Allocate<MachineBasicBlock>(Allocator))
              MachineBasicBlock(*this, bb);
 }
 
 /// Delete the given MachineBasicBlock.
 void
 MachineFunction::DeleteMachineBasicBlock(MachineBasicBlock *MBB) {
   assert(MBB->getParent() == this && "MBB parent mismatch!");
+  // Clean up any references to MBB in jump tables before deleting it
+  if (JumpTableInfo)
+    JumpTableInfo->RemoveMBBFromJumpTables(MBB);
   MBB->~MachineBasicBlock();
   BasicBlockRecycler.Deallocate(Allocator, MBB);
 }
 
 MachineMemOperand *MachineFunction::getMachineMemOperand(
     MachinePointerInfo PtrInfo, MachineMemOperand::Flags f, uint64_t s,
     Align base_alignment, const AAMDNodes &AAInfo, const MDNode *Ranges,
     SyncScope::ID SSID, AtomicOrdering Ordering,
bool MachineJumpTableInfo::ReplaceMBBInJumpTables(MachineBasicBlock *Old,
                                                   MachineBasicBlock *New) {
   assert(Old != New && "Not making a change?");
   bool MadeChange = false;
   for (size_t i = 0, e = JumpTables.size(); i != e; ++i)
     ReplaceMBBInJumpTable(i, Old, New);
   return MadeChange;
 }
 
+/// If MBB is present in any jump tables, remove it
+bool MachineJumpTableInfo::RemoveMBBFromJumpTables(MachineBasicBlock *MBB) {
+  bool MadeChange = false;
+  for (MachineJumpTableEntry &JTE : JumpTables) {
+    auto removeBeginItr = std::remove(JTE.MBBs.begin(), JTE.MBBs.end(), MBB);
+    MadeChange |= (removeBeginItr != JTE.MBBs.end());
+    JTE.MBBs.erase(removeBeginItr, JTE.MBBs.end());
+  }
+  return MadeChange;
+}
+
 /// If Old is a target of the jump tables, update the jump table to branch to
 /// New instead.
 bool MachineJumpTableInfo::ReplaceMBBInJumpTable(unsigned Idx,
                                                  MachineBasicBlock *Old,
                                                  MachineBasicBlock *New) {
   assert(Old != New && "Not making a change?");
   bool MadeChange = false;
   MachineJumpTableEntry &JTE = JumpTables[Idx];

llvm/test/CodeGen/PowerPC/unreachable-mbb-jtreference-elimination.ll

+; RUN: llc -O2 --verify-machineinstrs -stop-before=livevars \
+; RUN:   -mtriple=powerpc64le-unknown-unknown < %s | FileCheck %s
+
+
+define dso_local void @foo() #0 {
+; CHECK-LABEL: fixedStack:
+; CHECK-NEXT:  stack: []
+; CHECK-NEXT:  callSites: []
+; CHECK-NEXT:  debugValueSubstitutions: []
+; CHECK-NEXT:  constants: []
+; CHECK-NEXT:  machineFunctionInfo: {}
+; CHECK-NEXT:  jumpTable:
+; CHECK-NEXT:    kind:            label-difference32
+; CHECK-NEXT:    entries:
+; CHECK-NEXT:      - id:              0
+; CHECK-NEXT:        blocks:          [  ]
+; CHECK-NEXT:  body:             |
+; CHECK-NEXT:    bb.0.entry:
+; CHECK-NEXT:      successors: %bb.1(0x80000000)
+; CHECK:           B %bb.1
+; CHECK:         bb.1.next11:
+; CHECK-NEXT:      successors: %bb.2(0x80000000)
+; CHECK:           B %bb.2
+; CHECK:         bb.2.if.end139:
+entry:
+  br label %next11
+next11:                                           ; preds = %entry
+  br i1 false, label %if.then12, label %if.end139
+if.then12:                                        ; preds = %next11
+  br label %for.cond14
+for.cond14:                                       ; preds = %if.then12
+  switch i32 undef, label %sw.epilog [
+    i32 1, label %sw.bb
+    i32 4, label %sw.bb
+    i32 6, label %sw.bb
+    i32 7, label %sw.bb
+    i32 9, label %sw.bb
+    i32 12, label %sw.bb
+    i32 15, label %sw.bb
+    i32 16, label %sw.bb
+    i32 24, label %sw.bb
+    i32 0, label %sw.bb26
+    i32 2, label %sw.bb26
+    i32 3, label %sw.bb26
+    i32 8, label %sw.bb26
+    i32 10, label %sw.bb26
+    i32 11, label %sw.bb26
+    i32 13, label %sw.bb26
+    i32 17, label %sw.bb26
+    i32 18, label %sw.bb26
+    i32 20, label %sw.bb26
+    i32 19, label %sw.bb26
+    i32 21, label %sw.bb26
+    i32 22, label %sw.bb26
+    i32 23, label %sw.bb26
+    i32 25, label %sw.bb26
+    i32 27, label %sw.bb26
+    i32 28, label %sw.bb26
+    i32 26, label %sw.bb37
+    i32 29, label %sw.bb37
+    i32 30, label %sw.bb53
+  ]
+
+sw.bb:                                            ; preds = %for.cond14
+  unreachable
+sw.bb26:                                          ; preds = %for.cond14
+  unreachable
+sw.bb37:                                          ; preds = %for.cond14
+  unreachable
+sw.bb53:                                          ; preds = %for.cond14
+  unreachable
+sw.epilog:                                        ; preds = %for.cond14
+  unreachable
+if.end139:                                        ; preds = %next11
+  unreachable
+}
+attributes #0 = { noinline optnone }