From 086ec12628b64aed9bf0c46f97fa93f78c1f4abd Mon Sep 17 00:00:00 2001
From: Daniel Hoekwater <hoekwater@google.com>
Date: Thu, 3 Oct 2024 17:19:05 +0000
Subject: [PATCH 1/2] [CFIFixup] Clean up a redundant variable (NFC)

InsertMBB only ever refers to InsertPt->getParent(). Replace its usage
directly.
---
 llvm/lib/CodeGen/CFIFixup.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/CFIFixup.cpp b/llvm/lib/CodeGen/CFIFixup.cpp
index 61888a4266652..9bd10a99c5148 100644
--- a/llvm/lib/CodeGen/CFIFixup.cpp
+++ b/llvm/lib/CodeGen/CFIFixup.cpp
@@ -179,7 +179,6 @@ bool CFIFixup::runOnMachineFunction(MachineFunction &MF) {
   // `InsertPt` always points to the point in a preceding block where we have to
   // insert a `.cfi_remember_state`, in the case that the current block needs a
   // `.cfi_restore_state`.
-  MachineBasicBlock *InsertMBB = PrologueBlock;
   MachineBasicBlock::iterator InsertPt = PrologueEnd;
 
   assert(InsertPt != PrologueBlock->begin() &&
@@ -214,7 +213,7 @@ bool CFIFixup::runOnMachineFunction(MachineFunction &MF) {
       // stack frame.
       unsigned CFIIndex =
           MF.addFrameInst(MCCFIInstruction::createRememberState(nullptr));
-      BuildMI(*InsertMBB, InsertPt, DebugLoc(),
+      BuildMI(*InsertPt->getParent(), InsertPt, DebugLoc(),
               TII.get(TargetOpcode::CFI_INSTRUCTION))
           .addCFIIndex(CFIIndex);
       // Insert a `.cfi_restore_state` at the beginning of the current block.
@@ -223,7 +222,6 @@ bool CFIFixup::runOnMachineFunction(MachineFunction &MF) {
                          TII.get(TargetOpcode::CFI_INSTRUCTION))
                      .addCFIIndex(CFIIndex);
       ++InsertPt;
-      InsertMBB = &*CurrBB;
       Change = true;
     } else if ((Info.StrongNoFrameOnEntry || !Info.HasFrameOnEntry) &&
                HasFrame) {

From d4a9e3c211e8f22e01421a3a5c9ff9c9a764e883 Mon Sep 17 00:00:00 2001
From: Daniel Hoekwater <hoekwater@google.com>
Date: Thu, 3 Oct 2024 17:24:04 +0000
Subject: [PATCH 2/2] [CFIFixup] Factor remember/restore insertion into a
 helper function (NFC)

---
 llvm/lib/CodeGen/CFIFixup.cpp | 44 ++++++++++++++++++++++++-----------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/CodeGen/CFIFixup.cpp b/llvm/lib/CodeGen/CFIFixup.cpp
index 9bd10a99c5148..e881a62522f0c 100644
--- a/llvm/lib/CodeGen/CFIFixup.cpp
+++ b/llvm/lib/CodeGen/CFIFixup.cpp
@@ -116,6 +116,32 @@ findPrologueEnd(MachineFunction &MF, MachineBasicBlock::iterator &PrologueEnd) {
   return nullptr;
 }
 
+// Inserts a `.cfi_remember_state` instruction before PrologueEnd and a
+// `.cfi_restore_state` instruction before DstInsertPt. Returns an iterator
+// to the first instruction after the inserted `.cfi_restore_state` instruction.
+static MachineBasicBlock::iterator
+insertRememberRestorePair(MachineBasicBlock::iterator RememberInsertPt,
+                          MachineBasicBlock::iterator RestoreInsertPt) {
+  MachineBasicBlock *RememberMBB = RememberInsertPt->getParent();
+  MachineBasicBlock *RestoreMBB = RestoreInsertPt->getParent();
+  MachineFunction &MF = *RememberMBB->getParent();
+  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+
+  // Insert the `.cfi_remember_state` instruction.
+  unsigned CFIIndex =
+      MF.addFrameInst(MCCFIInstruction::createRememberState(nullptr));
+  BuildMI(*RememberMBB, RememberInsertPt, DebugLoc(),
+          TII.get(TargetOpcode::CFI_INSTRUCTION))
+      .addCFIIndex(CFIIndex);
+
+  // Insert the `.cfi_restore_state` instruction.
+  CFIIndex = MF.addFrameInst(MCCFIInstruction::createRestoreState(nullptr));
+  BuildMI(*RestoreMBB, RestoreInsertPt, DebugLoc(),
+          TII.get(TargetOpcode::CFI_INSTRUCTION))
+      .addCFIIndex(CFIIndex);
+  return RestoreInsertPt;
+}
+
 bool CFIFixup::runOnMachineFunction(MachineFunction &MF) {
   const TargetFrameLowering &TFL = *MF.getSubtarget().getFrameLowering();
   if (!TFL.enableCFIFixup(MF))
@@ -174,7 +200,6 @@ bool CFIFixup::runOnMachineFunction(MachineFunction &MF) {
   // Every block inherits the frame state (as recorded in the unwind tables)
   // of the previous block. If the intended frame state is different, insert
   // compensating CFI instructions.
-  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
   bool Change = false;
   // `InsertPt` always points to the point in a preceding block where we have to
   // insert a `.cfi_remember_state`, in the case that the current block needs a
@@ -209,19 +234,10 @@ bool CFIFixup::runOnMachineFunction(MachineFunction &MF) {
     if (!Info.StrongNoFrameOnEntry && Info.HasFrameOnEntry && !HasFrame) {
       // Reset to the "after prologue" state.
 
-      // Insert a `.cfi_remember_state` into the last block known to have a
-      // stack frame.
-      unsigned CFIIndex =
-          MF.addFrameInst(MCCFIInstruction::createRememberState(nullptr));
-      BuildMI(*InsertPt->getParent(), InsertPt, DebugLoc(),
-              TII.get(TargetOpcode::CFI_INSTRUCTION))
-          .addCFIIndex(CFIIndex);
-      // Insert a `.cfi_restore_state` at the beginning of the current block.
-      CFIIndex = MF.addFrameInst(MCCFIInstruction::createRestoreState(nullptr));
-      InsertPt = BuildMI(*CurrBB, CurrBB->begin(), DebugLoc(),
-                         TII.get(TargetOpcode::CFI_INSTRUCTION))
-                     .addCFIIndex(CFIIndex);
-      ++InsertPt;
+      // There's an earlier block known to have a stack frame. Insert a
+      // `.cfi_remember_state` instruction into that block and a
+      // `.cfi_restore_state` instruction at the beginning of the current block.
+      InsertPt = insertRememberRestorePair(InsertPt, CurrBB->begin());
       Change = true;
     } else if ((Info.StrongNoFrameOnEntry || !Info.HasFrameOnEntry) &&
                HasFrame) {