[MachineBasicBlock] Fix use after free in SplitCriticalEdge (#68786)

Remove use after free when attempting to update SlotIndexes in
MachineBasicBlock::SplitCriticalEdge.

Use MachineFunction delegate mechanism to capture target specific
manipulations of branch instructions and update SlotIndexes.

1 files changed, 27 insertions, 31 deletions

diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp
index 7d3d8b6..14d9bb29 100644
--- a/llvm/lib/CodeGen/MachineBasicBlock.cpp
+++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -1097,6 +1097,30 @@ static bool jumpTableHasOtherUses(const MachineFunction &MF,
   return false;
 }
 
+class SlotIndexUpdateDelegate : public MachineFunction::Delegate {
+private:
+  MachineFunction &MF;
+  SlotIndexes *Indexes;
+
+public:
+  SlotIndexUpdateDelegate(MachineFunction &MF, SlotIndexes *Indexes)
+      : MF(MF), Indexes(Indexes) {
+    MF.setDelegate(this);
+  }
+
+  ~SlotIndexUpdateDelegate() { MF.resetDelegate(this); }
+
+  void MF_HandleInsertion(MachineInstr &MI) override {
+    if (Indexes)
+      Indexes->insertMachineInstrInMaps(MI);
+  }
+
+  void MF_HandleRemoval(MachineInstr &MI) override {
+    if (Indexes)
+      Indexes->removeMachineInstrFromMaps(MI);
+  }
+};
+
 MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
     MachineBasicBlock *Succ, Pass &P,
     std::vector<SparseBitVector<>> *LiveInSets) {
@@ -1170,51 +1194,23 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
 
   ReplaceUsesOfBlockWith(Succ, NMBB);
 
-  // If updateTerminator() removes instructions, we need to remove them from
-  // SlotIndexes.
-  SmallVector<MachineInstr*, 4> Terminators;
-  if (Indexes) {
-    for (MachineInstr &MI :
-         llvm::make_range(getFirstInstrTerminator(), instr_end()))
-      Terminators.push_back(&MI);
-  }
-
   // Since we replaced all uses of Succ with NMBB, that should also be treated
   // as the fallthrough successor
   if (Succ == PrevFallthrough)
     PrevFallthrough = NMBB;
 
-  if (!ChangedIndirectJump)
+  if (!ChangedIndirectJump) {
+    SlotIndexUpdateDelegate SlotUpdater(*MF, Indexes);
     updateTerminator(PrevFallthrough);
-
-  if (Indexes) {
-    SmallVector<MachineInstr*, 4> NewTerminators;
-    for (MachineInstr &MI :
-         llvm::make_range(getFirstInstrTerminator(), instr_end()))
-      NewTerminators.push_back(&MI);
-
-    for (MachineInstr *Terminator : Terminators) {
-      if (!is_contained(NewTerminators, Terminator))
-        Indexes->removeMachineInstrFromMaps(*Terminator);
-    }
   }
 
   // Insert unconditional "jump Succ" instruction in NMBB if necessary.
   NMBB->addSuccessor(Succ);
   if (!NMBB->isLayoutSuccessor(Succ)) {
+    SlotIndexUpdateDelegate SlotUpdater(*MF, Indexes);
     SmallVector<MachineOperand, 4> Cond;
     const TargetInstrInfo *TII = getParent()->getSubtarget().getInstrInfo();
     TII->insertBranch(*NMBB, Succ, nullptr, Cond, DL);
-
-    if (Indexes) {
-      for (MachineInstr &MI : NMBB->instrs()) {
-        // Some instructions may have been moved to NMBB by updateTerminator(),
-        // so we first remove any instruction that already has an index.
-        if (Indexes->hasIndex(MI))
-          Indexes->removeMachineInstrFromMaps(MI);
-        Indexes->insertMachineInstrInMaps(MI);
-      }
-    }
   }
 
   // Fix PHI nodes in Succ so they refer to NMBB instead of this.