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

[Carl Ritson via llvm-commits]

https://github.com/perlfu created https://github.com/llvm/llvm-project/pull/68786

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

Add delegate mechanism for MachineBasicBlock to observe instructions added or removed when calling target specific functions for updating branches and update SlotIndexes appropriately.

>From 2f753b4ae96604ba76202b5cf011be1ec0adf49f Mon Sep 17 00:00:00 2001
From: Carl Ritson <carl.ritson at amd.com>
Date: Wed, 11 Oct 2023 19:10:50 +0900
Subject: [PATCH] [MachineBasicBlock] Fix use after free in SplitCriticalEdge

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

Add delegate mechanism for MachineBasicBlock to observe instructions
added or removed when calling target specific functions for updating
branches and update SlotIndexes appropriately.
---
 llvm/include/llvm/CodeGen/MachineBasicBlock.h | 37 +++++++++++
 llvm/lib/CodeGen/MachineBasicBlock.cpp        | 63 ++++++++++---------
 2 files changed, 69 insertions(+), 31 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachineBasicBlock.h b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
index 15c4fcd8399c181..4f9df77141debc8 100644
--- a/llvm/include/llvm/CodeGen/MachineBasicBlock.h
+++ b/llvm/include/llvm/CodeGen/MachineBasicBlock.h
@@ -14,6 +14,7 @@
 #define LLVM_CODEGEN_MACHINEBASICBLOCK_H
 
 #include "llvm/ADT/GraphTraits.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SparseBitVector.h"
 #include "llvm/ADT/ilist.h"
 #include "llvm/ADT/iterator_range.h"
@@ -106,6 +107,16 @@ class MachineBasicBlock
         : PhysReg(PhysReg), LaneMask(LaneMask) {}
   };
 
+  class Delegate {
+    virtual void anchor();
+
+  public:
+    virtual ~Delegate() = default;
+
+    virtual void MBB_NoteInsertion(MachineInstr &MI) = 0;
+    virtual void MBB_NoteRemoval(MachineInstr &MI) = 0;
+  };
+
 private:
   using Instructions = ilist<MachineInstr, ilist_sentinel_tracking<true>>;
 
@@ -233,6 +244,9 @@ class MachineBasicBlock
   /// Return a formatted string to identify this block and its parent function.
   std::string getFullName() const;
 
+  /// Delegates to inform of instruction addition and removal.
+  SmallPtrSet<Delegate *, 1> TheDelegates;
+
   /// Test whether this block is used as something other than the target
   /// of a terminator, exception-handling target, or jump table. This is
   /// either the result of an IR-level "blockaddress", or some form
@@ -1184,6 +1198,29 @@ class MachineBasicBlock
   /// MachineBranchProbabilityInfo class.
   BranchProbability getSuccProbability(const_succ_iterator Succ) const;
 
+  void resetDelegate(Delegate *delegate) {
+    assert(TheDelegates.count(delegate) &&
+           "Only an existing delegate can perform reset!");
+    TheDelegates.erase(delegate);
+  }
+
+  void addDelegate(Delegate *delegate) {
+    assert(delegate && !TheDelegates.count(delegate) &&
+           "Attempted to add null delegate, or to change it without "
+           "first resetting it!");
+    TheDelegates.insert(delegate);
+  }
+
+  void noteInsertion(MachineInstr &MI) {
+    for (auto *TheDelegate : TheDelegates)
+      TheDelegate->MBB_NoteInsertion(MI);
+  }
+
+  void noteRemoval(MachineInstr &MI) {
+    for (auto *TheDelegate : TheDelegates)
+      TheDelegate->MBB_NoteRemoval(MI);
+  }
+
 private:
   /// Return probability iterator corresponding to the I successor iterator.
   probability_iterator getProbabilityIterator(succ_iterator I);
diff --git a/llvm/lib/CodeGen/MachineBasicBlock.cpp b/llvm/lib/CodeGen/MachineBasicBlock.cpp
index 7d3d8b6fba1b7df..38dffc031dcb66d 100644
--- a/llvm/lib/CodeGen/MachineBasicBlock.cpp
+++ b/llvm/lib/CodeGen/MachineBasicBlock.cpp
@@ -156,12 +156,14 @@ void ilist_traits<MachineInstr>::addNodeToList(MachineInstr *N) {
   MachineFunction *MF = Parent->getParent();
   N->addRegOperandsToUseLists(MF->getRegInfo());
   MF->handleInsertion(*N);
+  Parent->noteInsertion(*N);
 }
 
 /// When we remove an instruction from a basic block list, we update its parent
 /// pointer and remove its operands from reg use/def lists if appropriate.
 void ilist_traits<MachineInstr>::removeNodeFromList(MachineInstr *N) {
   assert(N->getParent() && "machine instruction not in a basic block");
+  N->getParent()->noteRemoval(*N);
 
   // Remove from the use/def lists.
   if (MachineFunction *MF = N->getMF()) {
@@ -1097,6 +1099,31 @@ static bool jumpTableHasOtherUses(const MachineFunction &MF,
   return false;
 }
 
+class MBBSplitCriticalEdgeDelegate : public MachineBasicBlock::Delegate {
+private:
+  MachineBasicBlock *MBB;
+  SlotIndexes *Indexes;
+
+public:
+  MBBSplitCriticalEdgeDelegate(MachineBasicBlock *Block,
+                               SlotIndexes *IndexesToUpdate)
+      : MBB(Block), Indexes(IndexesToUpdate) {
+    MBB->addDelegate(this);
+  }
+
+  ~MBBSplitCriticalEdgeDelegate() { MBB->resetDelegate(this); }
+
+  void MBB_NoteInsertion(MachineInstr &MI) override {
+    if (Indexes)
+      Indexes->insertMachineInstrInMaps(MI);
+  }
+
+  void MBB_NoteRemoval(MachineInstr &MI) override {
+    if (Indexes)
+      Indexes->removeMachineInstrFromMaps(MI);
+  }
+};
+
 MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
     MachineBasicBlock *Succ, Pass &P,
     std::vector<SparseBitVector<>> *LiveInSets) {
@@ -1170,51 +1197,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) {
+    MBBSplitCriticalEdgeDelegate SlotUpdater(this, 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)) {
+    MBBSplitCriticalEdgeDelegate SlotUpdater(NMBB, 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.
@@ -1743,3 +1742,5 @@ bool MachineBasicBlock::sizeWithoutDebugLargerThan(unsigned Limit) const {
 const MBBSectionID MBBSectionID::ColdSectionID(MBBSectionID::SectionType::Cold);
 const MBBSectionID
     MBBSectionID::ExceptionSectionID(MBBSectionID::SectionType::Exception);
+
+void MachineBasicBlock::Delegate::anchor() {}