[AArch64][PAC] Fix creating check instructions for BBs without an epilog (#92508)

`AArch64PAuth::checkAuthenticatedRegister()` splits the basic block
containing the tail call instruction to add check instructions, assuming
at least one more instruction before the call. This assumption is
incorrect in cases where some execution paths lead to the termination
block without creating the stack frame. This patch rearranges the
creation of the checks so that the prior splitting is not required.

Author: igorkudrin

llvm/lib/Target/AArch64/AArch64PointerAuth.cpp

@@ -241,6 +241,14 @@ void llvm::AArch64PAuth::checkAuthenticatedRegister(
   const AArch64InstrInfo *TII = Subtarget.getInstrInfo();
   DebugLoc DL = MBBI->getDebugLoc();
 
+  // All terminator instructions should be grouped at the end of the machine
+  // basic block, with no non-terminator instructions between them. Depending on
+  // the method requested, we will insert some regular instructions, maybe
+  // followed by a conditional branch instruction, which is a terminator, before
+  // MBBI. Thus, MBBI is expected to be the first terminator of its MBB.
+  assert(MBBI->isTerminator() && MBBI == MBB.getFirstTerminator() &&
+         "MBBI should be the first terminator in MBB");
+
   // First, handle the methods not requiring creating extra MBBs.
   switch (Method) {
   default:
@@ -257,33 +265,24 @@ void llvm::AArch64PAuth::checkAuthenticatedRegister(
 
   // Control flow has to be changed, so arrange new MBBs.
 
-  // At now, at least an AUT* instruction is expected before MBBI
-  assert(MBBI != MBB.begin() &&
-         "Cannot insert the check at the very beginning of MBB");
-  // The block to insert check into.
-  MachineBasicBlock *CheckBlock = &MBB;
-  // The remaining part of the original MBB that is executed on success.
-  MachineBasicBlock *SuccessBlock = MBB.splitAt(*std::prev(MBBI));
-
   // The block that explicitly generates a break-point exception on failure.
   MachineBasicBlock *BreakBlock =
       MF.CreateMachineBasicBlock(MBB.getBasicBlock());
   MF.push_back(BreakBlock);
-  MBB.splitSuccessor(SuccessBlock, BreakBlock);
+  MBB.addSuccessor(BreakBlock);
 
-  assert(CheckBlock->getFallThrough() == SuccessBlock);
   BuildMI(BreakBlock, DL, TII->get(AArch64::BRK)).addImm(BrkImm);
 
   switch (Method) {
   case AuthCheckMethod::None:
   case AuthCheckMethod::DummyLoad:
     llvm_unreachable("Should be handled above");
   case AuthCheckMethod::HighBitsNoTBI:
-    BuildMI(CheckBlock, DL, TII->get(AArch64::EORXrs), TmpReg)
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::EORXrs), TmpReg)
         .addReg(AuthenticatedReg)
         .addReg(AuthenticatedReg)
         .addImm(1);
-    BuildMI(CheckBlock, DL, TII->get(AArch64::TBNZX))
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::TBNZX))
         .addReg(TmpReg)
         .addImm(62)
         .addMBB(BreakBlock);
@@ -292,16 +291,16 @@ void llvm::AArch64PAuth::checkAuthenticatedRegister(
     assert(AuthenticatedReg == AArch64::LR &&
            "XPACHint mode is only compatible with checking the LR register");
     assert(UseIKey && "XPACHint mode is only compatible with I-keys");
-    BuildMI(CheckBlock, DL, TII->get(AArch64::ORRXrs), TmpReg)
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::ORRXrs), TmpReg)
         .addReg(AArch64::XZR)
         .addReg(AArch64::LR)
         .addImm(0);
-    BuildMI(CheckBlock, DL, TII->get(AArch64::XPACLRI));
-    BuildMI(CheckBlock, DL, TII->get(AArch64::SUBSXrs), AArch64::XZR)
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::XPACLRI));
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::SUBSXrs), AArch64::XZR)
         .addReg(TmpReg)
         .addReg(AArch64::LR)
         .addImm(0);
-    BuildMI(CheckBlock, DL, TII->get(AArch64::Bcc))
+    BuildMI(MBB, MBBI, DL, TII->get(AArch64::Bcc))
         .addImm(AArch64CC::NE)
         .addMBB(BreakBlock);
     return;

llvm/test/CodeGen/AArch64/sign-return-address-tailcall.ll

@@ -129,4 +129,36 @@ define i32 @tailcall_ib_key() "sign-return-address"="all" "sign-return-address-k
   ret i32 %call
 }
 
+define i32 @tailcall_two_branches(i1 %0) "sign-return-address"="all" {
+; COMMON-LABEL:    tailcall_two_branches:
+; COMMON:            tbz w0, #0, .[[ELSE:LBB[_0-9]+]]
+; COMMON:            str x30, [sp, #-16]!
+; COMMON:            bl callee2
+; COMMON:            ldr x30, [sp], #16
+; COMMON-NEXT:       [[AUTIASP]]
+; COMMON-NEXT:     .[[ELSE]]:
+
+; LDR-NEXT:          ldr w16, [x30]
+;
+; BITS-NOTBI-NEXT:   eor x16, x30, x30, lsl #1
+; BITS-NOTBI-NEXT:   tbnz x16, #62, .[[FAIL:LBB[_0-9]+]]
+;
+; XPAC-NEXT:         mov x16, x30
+; XPAC-NEXT:         [[XPACLRI]]
+; XPAC-NEXT:         cmp x16, x30
+; XPAC-NEXT:         b.ne .[[FAIL:LBB[_0-9]+]]
+;
+; COMMON-NEXT:       b callee
+; BRK-NEXT:        .[[FAIL]]:
+; BRK-NEXT:          brk #0xc470
+  br i1 %0, label %2, label %3
+2:
+  call void @callee2()
+  br label %3
+3:
+  %call = tail call i32 @callee()
+  ret i32 %call
+}
+
 declare i32 @callee()
+declare void @callee2()