[llvm] [AArch64][PAC] Support BLRA* instructions in SLS Hardening pass (PR #97605)

[Anatoly Trosinenko via llvm-commits]

================
@@ -139,64 +265,59 @@ bool SLSHardeningInserter::hardenReturnsAndBRs(MachineModuleInfo &MMI,
   return Modified;
 }
 
-static const unsigned NumPermittedRegs = 29;
-static const struct ThunkNameAndReg {
-  const char* Name;
-  Register Reg;
-} SLSBLRThunks[NumPermittedRegs] = {
-    {"__llvm_slsblr_thunk_x0", AArch64::X0},
-    {"__llvm_slsblr_thunk_x1", AArch64::X1},
-    {"__llvm_slsblr_thunk_x2", AArch64::X2},
-    {"__llvm_slsblr_thunk_x3", AArch64::X3},
-    {"__llvm_slsblr_thunk_x4", AArch64::X4},
-    {"__llvm_slsblr_thunk_x5", AArch64::X5},
-    {"__llvm_slsblr_thunk_x6", AArch64::X6},
-    {"__llvm_slsblr_thunk_x7", AArch64::X7},
-    {"__llvm_slsblr_thunk_x8", AArch64::X8},
-    {"__llvm_slsblr_thunk_x9", AArch64::X9},
-    {"__llvm_slsblr_thunk_x10", AArch64::X10},
-    {"__llvm_slsblr_thunk_x11", AArch64::X11},
-    {"__llvm_slsblr_thunk_x12", AArch64::X12},
-    {"__llvm_slsblr_thunk_x13", AArch64::X13},
-    {"__llvm_slsblr_thunk_x14", AArch64::X14},
-    {"__llvm_slsblr_thunk_x15", AArch64::X15},
-    // X16 and X17 are deliberately missing, as the mitigation requires those
-    // register to not be used in BLR. See comment in ConvertBLRToBL for more
-    // details.
-    {"__llvm_slsblr_thunk_x18", AArch64::X18},
-    {"__llvm_slsblr_thunk_x19", AArch64::X19},
-    {"__llvm_slsblr_thunk_x20", AArch64::X20},
-    {"__llvm_slsblr_thunk_x21", AArch64::X21},
-    {"__llvm_slsblr_thunk_x22", AArch64::X22},
-    {"__llvm_slsblr_thunk_x23", AArch64::X23},
-    {"__llvm_slsblr_thunk_x24", AArch64::X24},
-    {"__llvm_slsblr_thunk_x25", AArch64::X25},
-    {"__llvm_slsblr_thunk_x26", AArch64::X26},
-    {"__llvm_slsblr_thunk_x27", AArch64::X27},
-    {"__llvm_slsblr_thunk_x28", AArch64::X28},
-    {"__llvm_slsblr_thunk_x29", AArch64::FP},
-    // X30 is deliberately missing, for similar reasons as X16 and X17 are
-    // missing.
-    {"__llvm_slsblr_thunk_x31", AArch64::XZR},
-};
+// Currently, the longest possible thunk name is
+//   __llvm_slsblr_thunk_aa_xNN_xMM
+// which is 31 characters (without the '\0' character).
+static SmallString<32> createThunkName(const ThunkKind &Kind, Register Xn,
+                                       Register Xm) {
+  unsigned N = ThunksSet::indexOfXReg(Xn);
+  if (!Kind.HasXmOperand)
+    return formatv("{0}{1}x{2}", CommonNamePrefix, Kind.NameInfix, N);
+
+  unsigned M = ThunksSet::indexOfXReg(Xm);
+  return formatv("{0}{1}x{2}_x{3}", CommonNamePrefix, Kind.NameInfix, N, M);
+}
 
-unsigned getThunkIndex(Register Reg) {
-  for (unsigned I = 0; I < NumPermittedRegs; ++I)
-    if (SLSBLRThunks[I].Reg == Reg)
-      return I;
-  llvm_unreachable("Unexpected register");
+static const ThunkKind &parseThunkName(StringRef ThunkName, Register &Xn,
+                                       Register &Xm) {
+  assert(ThunkName.starts_with(CommonNamePrefix) &&
+         "Should be filtered out by ThunkInserter");
+  // Thunk name suffix, such as "x1" or "aa_x2_x3".
+  StringRef NameSuffix = ThunkName.drop_front(CommonNamePrefix.size());
+
+  // Parse thunk kind based on thunk name infix.
+  const ThunkKind &Kind = *StringSwitch<const ThunkKind *>(NameSuffix)
+                               .StartsWith("aa_", &ThunkKind::BRAA)
+                               .StartsWith("ab_", &ThunkKind::BRAB)
+                               .StartsWith("aaz_", &ThunkKind::BRAAZ)
+                               .StartsWith("abz_", &ThunkKind::BRABZ)
+                               .Default(&ThunkKind::BR);
+
+  auto ParseRegName = [](StringRef Name) {
+    assert(Name.starts_with("x") && "xN register name expected");
+    unsigned N;
+    bool Fail = Name.drop_front(1).getAsInteger(/*Radix=*/10, N);
+    assert(!Fail && N < 32 && "Unexpected register");
----------------
atrosinenko wrote:

Fixed, thanks!

https://github.com/llvm/llvm-project/pull/97605