[AArch64][PAC] Use enum to describe LR signing condition (NFC) (#168548)

Express the condition of signing the return address in a function using
an `enum class` instead of a pair of `bool`s. Define `enum class
SignReturnAddress` with the values corresponding to the three possible
modes that can be requested via "sign-return-address" function
attribute.

Previously, there were two overloads of `shouldSignReturnAddress`
accepting either `const MachineFunction &` or `bool` argument. Due to
pointer-to-bool conversion, when `shouldSignReturnAddress` was
incorrectly called with `const MachineFunction *` argument, the latter
overload was used instead of reporting a compile-time error.

Author: atrosinenko

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -973,8 +973,7 @@ bool AArch64FrameLowering::shouldSignReturnAddressEverywhere(
   if (MF.getTarget().getMCAsmInfo()->usesWindowsCFI())
     return false;
   const AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
-  bool SignReturnAddressAll = AFI->shouldSignReturnAddress(/*SpillsLR=*/false);
-  return SignReturnAddressAll;
+  return AFI->getSignReturnAddressCondition() == SignReturnAddress::All;
 }
 
 // Given a load or a store instruction, generate an appropriate unwinding SEH

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -9796,8 +9796,8 @@ outliningCandidatesSigningScopeConsensus(const outliner::Candidate &a,
   const auto &MFIa = a.getMF()->getInfo<AArch64FunctionInfo>();
   const auto &MFIb = b.getMF()->getInfo<AArch64FunctionInfo>();
 
-  return MFIa->shouldSignReturnAddress(false) == MFIb->shouldSignReturnAddress(false) &&
-         MFIa->shouldSignReturnAddress(true) == MFIb->shouldSignReturnAddress(true);
+  return MFIa->getSignReturnAddressCondition() ==
+         MFIb->getSignReturnAddressCondition();
 }
 
 static bool
@@ -9888,10 +9888,11 @@ AArch64InstrInfo::getOutliningCandidateInfo(
   // Performing a tail call may require extra checks when PAuth is enabled.
   // If PAuth is disabled, set it to zero for uniformity.
   unsigned NumBytesToCheckLRInTCEpilogue = 0;
-  if (RepeatedSequenceLocs[0]
-          .getMF()
-          ->getInfo<AArch64FunctionInfo>()
-          ->shouldSignReturnAddress(true)) {
+  const auto RASignCondition = RepeatedSequenceLocs[0]
+                                   .getMF()
+                                   ->getInfo<AArch64FunctionInfo>()
+                                   ->getSignReturnAddressCondition();
+  if (RASignCondition != SignReturnAddress::None) {
     // One PAC and one AUT instructions
     NumBytesToCreateFrame += 8;
 
@@ -10695,7 +10696,9 @@ void AArch64InstrInfo::buildOutlinedFrame(
     Et = MBB.insert(Et, LDRXpost);
   }
 
-  bool ShouldSignReturnAddr = FI->shouldSignReturnAddress(!IsLeafFunction);
+  auto RASignCondition = FI->getSignReturnAddressCondition();
+  bool ShouldSignReturnAddr = AArch64FunctionInfo::shouldSignReturnAddress(
+      RASignCondition, !IsLeafFunction);
 
   // If this is a tail call outlined function, then there's already a return.
   if (OF.FrameConstructionID == MachineOutlinerTailCall ||

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp

@@ -16,6 +16,7 @@
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
+#include "llvm/ADT/StringSwitch.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
@@ -63,24 +64,21 @@ void AArch64FunctionInfo::initializeBaseYamlFields(
     setHasStreamingModeChanges(*YamlMFI.HasStreamingModeChanges);
 }
 
-static std::pair<bool, bool> GetSignReturnAddress(const Function &F) {
+static SignReturnAddress GetSignReturnAddress(const Function &F) {
   if (F.hasFnAttribute("ptrauth-returns"))
-    return {true, false}; // non-leaf
+    return SignReturnAddress::NonLeaf;
+
   // The function should be signed in the following situations:
   // - sign-return-address=all
   // - sign-return-address=non-leaf and the functions spills the LR
   if (!F.hasFnAttribute("sign-return-address"))
-    return {false, false};
+    return SignReturnAddress::None;
 
   StringRef Scope = F.getFnAttribute("sign-return-address").getValueAsString();
-  if (Scope == "none")
-    return {false, false};
-
-  if (Scope == "all")
-    return {true, true};
-
-  assert(Scope == "non-leaf");
-  return {true, false};
+  return StringSwitch<SignReturnAddress>(Scope)
+      .Case("none", SignReturnAddress::None)
+      .Case("non-leaf", SignReturnAddress::NonLeaf)
+      .Case("all", SignReturnAddress::All);
 }
 
 static bool ShouldSignWithBKey(const Function &F, const AArch64Subtarget &STI) {
@@ -116,7 +114,7 @@ AArch64FunctionInfo::AArch64FunctionInfo(const Function &F,
   // HasRedZone here.
   if (F.hasFnAttribute(Attribute::NoRedZone))
     HasRedZone = false;
-  std::tie(SignReturnAddress, SignReturnAddressAll) = GetSignReturnAddress(F);
+  SignCondition = GetSignReturnAddress(F);
   SignWithBKey = ShouldSignWithBKey(F, *STI);
   HasELFSignedGOT = hasELFSignedGOTHelper(F, STI);
   // TODO: skip functions that have no instrumented allocas for optimization
@@ -169,23 +167,27 @@ MachineFunctionInfo *AArch64FunctionInfo::clone(
   return DestMF.cloneInfo<AArch64FunctionInfo>(*this);
 }
 
-bool AArch64FunctionInfo::shouldSignReturnAddress(bool SpillsLR) const {
-  if (!SignReturnAddress)
-    return false;
-  if (SignReturnAddressAll)
-    return true;
-  return SpillsLR;
-}
-
 static bool isLRSpilled(const MachineFunction &MF) {
   return llvm::any_of(
       MF.getFrameInfo().getCalleeSavedInfo(),
       [](const auto &Info) { return Info.getReg() == AArch64::LR; });
 }
 
+bool AArch64FunctionInfo::shouldSignReturnAddress(SignReturnAddress Condition,
+                                                  bool IsLRSpilled) {
+  switch (Condition) {
+  case SignReturnAddress::None:
+    return false;
+  case SignReturnAddress::NonLeaf:
+    return IsLRSpilled;
+  case SignReturnAddress::All:
+    return true;
+  }
+}
+
 bool AArch64FunctionInfo::shouldSignReturnAddress(
     const MachineFunction &MF) const {
-  return shouldSignReturnAddress(isLRSpilled(MF));
+  return shouldSignReturnAddress(SignCondition, isLRSpilled(MF));
 }
 
 bool AArch64FunctionInfo::needsShadowCallStackPrologueEpilogue(

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -42,6 +42,15 @@ struct TPIDR2Object {
   unsigned Uses = 0;
 };
 
+/// Condition of signing the return address in a function.
+///
+/// Corresponds to possible values of "sign-return-address" function attribute.
+enum class SignReturnAddress {
+  None,
+  NonLeaf,
+  All,
+};
+
 /// AArch64FunctionInfo - This class is derived from MachineFunctionInfo and
 /// contains private AArch64-specific information for each MachineFunction.
 class AArch64FunctionInfo final : public MachineFunctionInfo {
@@ -170,13 +179,8 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
   // CalleeSavedStackSize) to the address of the frame record.
   int CalleeSaveBaseToFrameRecordOffset = 0;
 
-  /// SignReturnAddress is true if PAC-RET is enabled for the function with
-  /// defaults being sign non-leaf functions only, with the B key.
-  bool SignReturnAddress = false;
-
-  /// SignReturnAddressAll modifies the default PAC-RET mode to signing leaf
-  /// functions as well.
-  bool SignReturnAddressAll = false;
+  /// SignCondition controls when PAC-RET protection should be used.
+  SignReturnAddress SignCondition = SignReturnAddress::None;
 
   /// SignWithBKey modifies the default PAC-RET mode to signing with the B key.
   bool SignWithBKey = false;
@@ -591,8 +595,14 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
     CalleeSaveBaseToFrameRecordOffset = Offset;
   }
 
+  static bool shouldSignReturnAddress(SignReturnAddress Condition,
+                                      bool IsLRSpilled);
+
   bool shouldSignReturnAddress(const MachineFunction &MF) const;
-  bool shouldSignReturnAddress(bool SpillsLR) const;
+
+  SignReturnAddress getSignReturnAddressCondition() const {
+    return SignCondition;
+  }
 
   bool needsShadowCallStackPrologueEpilogue(MachineFunction &MF) const;