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];