1 files changed, 38 insertions, 2 deletions
diff --git a/gnu/llvm/lib/CodeGen/PrologEpilogInserter.cpp b/gnu/llvm/lib/CodeGen/PrologEpilogInserter.cpp
index a8d8ad8ac7d..529fa7406a5 100644
--- a/gnu/llvm/lib/CodeGen/PrologEpilogInserter.cpp
+++ b/gnu/llvm/lib/CodeGen/PrologEpilogInserter.cpp
@@ -175,6 +175,10 @@ bool PEI::runOnMachineFunction(MachineFunction &Fn) {
   const TargetRegisterInfo *TRI = Fn.getSubtarget().getRegisterInfo();
   const TargetFrameLowering *TFI = Fn.getSubtarget().getFrameLowering();
 
+  // Set Return Protector in the frame
+  if (F.hasFnAttribute("ret-protector"))
+    Fn.getFrameInfo().setReturnProtector(true);
+
   RS = TRI->requiresRegisterScavenging(Fn) ? new RegScavenger() : nullptr;
   FrameIndexVirtualScavenging = TRI->requiresFrameIndexScavenging(Fn);
   FrameIndexEliminationScavenging = (RS && !FrameIndexVirtualScavenging) ||
@@ -209,6 +213,19 @@ bool PEI::runOnMachineFunction(MachineFunction &Fn) {
   if (!F.hasFnAttribute(Attribute::Naked))
     insertPrologEpilogCode(Fn);
 
+  // Add Return Protectors if using them
+  if (Fn.getFrameInfo().hasReturnProtector()) {
+    std::vector<MachineBasicBlock *> ReturnBlocks;
+    bool insertedGuard = false;
+    for (auto &MBB: Fn)
+      if (MBB.isReturnBlock())
+        ReturnBlocks.push_back(&MBB);
+    for (MachineBasicBlock *MBB: ReturnBlocks)
+      insertedGuard |= TFI->insertReturnProtectorEpilogue(Fn, *MBB);
+    if (insertedGuard)
+      TFI->insertReturnProtectorPrologue(Fn, Fn.front());
+  }
+
   // Replace all MO_FrameIndex operands with physical register references
   // and actual offsets.
   //
@@ -293,7 +310,8 @@ void PEI::calculateCallFrameInfo(MachineFunction &Fn) {
 /// Compute the sets of entry and return blocks for saving and restoring
 /// callee-saved registers, and placing prolog and epilog code.
 void PEI::calculateSaveRestoreBlocks(MachineFunction &Fn) {
-  const MachineFrameInfo &MFI = Fn.getFrameInfo();
+  MachineFrameInfo &MFI = Fn.getFrameInfo();
+  const TargetFrameLowering *TFI = Fn.getSubtarget().getFrameLowering();
 
   // Even when we do not change any CSR, we still want to insert the
   // prologue and epilogue of the function.
@@ -309,7 +327,18 @@ void PEI::calculateSaveRestoreBlocks(MachineFunction &Fn) {
     // epilogue.
     if (!RestoreBlock->succ_empty() || RestoreBlock->isReturnBlock())
       RestoreBlocks.push_back(RestoreBlock);
-    return;
+
+    // If we are adding return protectors ensure we can find a free register
+    if (MFI.hasReturnProtector() &&
+        !TFI->determineReturnProtectorTempRegister(Fn, SaveBlocks, RestoreBlocks)) {
+      // Shrinkwrapping will prevent finding a free register
+      SaveBlocks.clear();
+      RestoreBlocks.clear();
+      MFI.setSavePoint(nullptr);
+      MFI.setRestorePoint(nullptr);
+    } else {
+      return;
+    }
   }
 
   // Save refs to entry and return blocks.
@@ -320,6 +349,9 @@ void PEI::calculateSaveRestoreBlocks(MachineFunction &Fn) {
     if (MBB.isReturnBlock())
       RestoreBlocks.push_back(&MBB);
   }
+
+  if (MFI.hasReturnProtector())
+    TFI->determineReturnProtectorTempRegister(Fn, SaveBlocks, RestoreBlocks);
 }
 
 static void assignCalleeSavedSpillSlots(MachineFunction &F,
@@ -341,6 +373,10 @@ static void assignCalleeSavedSpillSlots(MachineFunction &F,
 
   const TargetFrameLowering *TFI = F.getSubtarget().getFrameLowering();
   MachineFrameInfo &MFI = F.getFrameInfo();
+
+  if (MFI.hasReturnProtectorTempRegister())
+      CSI.push_back(CalleeSavedInfo(MFI.getReturnProtectorTempRegister()));
+
   if (!TFI->assignCalleeSavedSpillSlots(F, RegInfo, CSI)) {
     // If target doesn't implement this, use generic code.