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

[Anatoly Trosinenko via llvm-commits]

================
@@ -32,17 +34,107 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
-static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+// Common name prefix of all thunks generated by this pass.
+//
+// The generic form is
+// __llvm_slsblr_thunk_xN            for BLR thunks
+// __llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
+// __llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks
+static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";
 
 namespace {
 
-// Set of inserted thunks: bitmask with bits corresponding to
-// indexes in SLSBLRThunks array.
-typedef uint32_t ThunksSet;
+struct ThunkKind {
+  enum ThunkKindId {
+    ThunkBR,
+    ThunkBRAA,
+    ThunkBRAB,
+    ThunkBRAAZ,
+    ThunkBRABZ,
+  };
+
+  ThunkKindId Id;
+  StringRef NameInfix;
+  bool HasXmOperand;
+  bool NeedsPAuth;
+
+  // Opcode to perform indirect jump from inside the thunk.
+  unsigned BROpcode;
+
+  static const ThunkKind BR;
+  static const ThunkKind BRAA;
+  static const ThunkKind BRAB;
+  static const ThunkKind BRAAZ;
+  static const ThunkKind BRABZ;
+};
+
+// Set of inserted thunks.
+class ThunksSet {
+public:
+  static constexpr unsigned NumXRegisters = 32;
+
+  // Given Xn register, returns n.
+  static unsigned indexOfXReg(Register Xn);
+  // Given n, returns Xn register.
+  static Register xRegByIndex(unsigned N);
+
+  ThunksSet &operator|=(const ThunksSet &Other) {
+    BLRThunks |= Other.BLRThunks;
+    BLRAAZThunks |= Other.BLRAAZThunks;
+    BLRABZThunks |= Other.BLRABZThunks;
+    for (unsigned I = 0; I < NumXRegisters; ++I)
+      BLRAAThunks[I] |= Other.BLRAAThunks[I];
+    for (unsigned I = 0; I < NumXRegisters; ++I)
+      BLRABThunks[I] |= Other.BLRABThunks[I];
+
+    return *this;
+  }
+
+  bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+    reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);
+    return getBitmask(Kind, Xm) & XnBit;
+  }
+
+  void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+    reg_bitmask_t XnBit = reg_bitmask_t(1) << indexOfXReg(Xn);
+    getBitmask(Kind, Xm) |= XnBit;
+  }
+
+private:
+  typedef uint32_t reg_bitmask_t;
+  static_assert(NumXRegisters <= sizeof(reg_bitmask_t) * CHAR_BIT,
----------------
atrosinenko wrote:

Added `<climits>` and `<tuple>` includes, thank you.

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