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

[Daniil Kovalev via llvm-commits]

================
@@ -32,17 +34,103 @@ 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) {
+    uint32_t XnBit = 1u << indexOfXReg(Xn);
+    return getBitmask(Kind, Xm) & XnBit;
+  }
+
+  void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+    uint32_t XnBit = 1u << indexOfXReg(Xn);
+    getBitmask(Kind, Xm) |= XnBit;
+  }
+
+private:
+  // Bitmasks representing operands used, with n-th bit corresponding to Xn
+  // register operand. If the instruction has a second operand (Xm), an array
+  // of bitmasks is used, indexed by m.
+  // Indexes corresponding to the forbidden x16, x17 and x30 registers are
+  // always unset, for simplicity there are no holes.
----------------
kovdan01 wrote:

Nit: a static assertion like this could be added:

```
static_assert(sizeof(uint32_t) * CHAR_BIT == NumXRegisters,
              "X-registers number must match the bitmask width");
```

Feel free to ignore - I'm not sure if it's nice and actually gives that much value

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