[llvm] [AArch64][PAC] Lower jump-tables using hardened pseudo. (PR #97666)

[Ahmed Bougacha via llvm-commits]

================
@@ -1310,6 +1312,138 @@ void AArch64AsmPrinter::LowerJumpTableDest(llvm::MCStreamer &OutStreamer,
                                   .addImm(Size == 4 ? 0 : 2));
 }
 
+void AArch64AsmPrinter::LowerHardenedBRJumpTable(const MachineInstr &MI) {
+  unsigned InstsEmitted = 0;
+
+  const MachineJumpTableInfo *MJTI = MF->getJumpTableInfo();
+  assert(MJTI && "Can't lower jump-table dispatch without JTI");
+
+  const std::vector<MachineJumpTableEntry> &JTs = MJTI->getJumpTables();
+  assert(!JTs.empty() && "Invalid JT index for jump-table dispatch");
+
+  // Emit:
+  //     mov x17, #<size of table>     ; depending on table size, with MOVKs
+  //     cmp x16, x17                  ; or #imm if table size fits in 12-bit
+  //     csel x16, x16, xzr, ls        ; check for index overflow
+  //
+  //     adrp x17, Ltable at PAGE         ; materialize table address
+  //     add x17, Ltable at PAGEOFF
+  //     ldrsw x16, [x17, x16, lsl #2] ; load table entry
+  //
+  //   Lanchor:
+  //     adr x17, Lanchor              ; compute target address
+  //     add x16, x17, x16
+  //     br x16                        ; branch to target
+
+  MachineOperand JTOp = MI.getOperand(0);
+
+  unsigned JTI = JTOp.getIndex();
+  assert(!AArch64FI->getJumpTableEntryPCRelSymbol(JTI) &&
+         "unsupported compressed jump table");
+
+  const uint64_t NumTableEntries = JTs[JTI].MBBs.size();
+
+  // cmp only supports a 12-bit immediate.  If we need more, materialize the
+  // immediate, using x17 as a scratch register.
+  uint64_t MaxTableEntry = NumTableEntries - 1;
+  if (isUInt<12>(MaxTableEntry)) {
+    EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::SUBSXri)
+                                     .addReg(AArch64::XZR)
+                                     .addReg(AArch64::X16)
+                                     .addImm(MaxTableEntry)
+                                     .addImm(0));
+    ++InstsEmitted;
+  } else {
+    EmitToStreamer(*OutStreamer,
+                   MCInstBuilder(AArch64::MOVZXi)
+                       .addReg(AArch64::X17)
+                       .addImm(static_cast<uint16_t>(MaxTableEntry))
+                       .addImm(0));
+    ++InstsEmitted;
+    // It's sad that we have to manually materialize instructions, but we can't
+    // trivially reuse the main pseudo expansion logic.
----------------
ahmedbougacha wrote:

Yeah, there's the high-level problem of "we have to materialize the immediate within this sequence to preserve its integrity, so we can't easily use the usual MOVimm pseudos".
But there's also the lower-level problem of "MOVimm pseudo expansion knows a lot of tricks that aren't worth reimplementing here".  IIRC this original implementation predates the nice and tidy expandMOVImm helper, which exposed that generic pseudo expansion.  So we could switch this to use that here.

But looking at this again, I don't think we'd want to do that anyway, in part because the table size is never going to be that large anyway (12 bits might even be always enough TBH, but I can think of pathological codebases where that might be false).  In part because the general MOVImm expansion is somewhat unpredictable in ways that would be annoying here (e.g., the pseudo size), and we already ended up with a similar "materialize an immediate into x17 in a predictable way" for other ptrauth pseudos.  So I think I'll just switch this to using that.

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