[AArch64][PAC] Fix creating check instructions for BBs without an epilog

llvm/lib/Target/AArch64/AArch64PointerAuth.cpp

@@ -231,7 +231,7 @@ MachineMemOperand *createCheckMemOperand(MachineFunction &MF,
 
 } // namespace
 
-MachineBasicBlock &llvm::AArch64PAuth::checkAuthenticatedRegister(
+void llvm::AArch64PAuth::checkAuthenticatedRegister(
     MachineBasicBlock::iterator MBBI, AuthCheckMethod Method,
     Register AuthenticatedReg, Register TmpReg, bool UseIKey, unsigned BrkImm) {
 
@@ -241,70 +241,69 @@ MachineBasicBlock &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:
     break;
   case AuthCheckMethod::None:
-    return MBB;
+    return;
   case AuthCheckMethod::DummyLoad:
     BuildMI(MBB, MBBI, DL, TII->get(AArch64::LDRWui), getWRegFromXReg(TmpReg))
         .addReg(AuthenticatedReg)
         .addImm(0)
         .addMemOperand(createCheckMemOperand(MF, Subtarget));
-    return MBB;
+    return;
   }
 
   // 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);
-    return *SuccessBlock;
+    return;
   case AuthCheckMethod::XPACHint:
     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 *SuccessBlock;
+    return;
   }
   llvm_unreachable("Unknown AuthCheckMethod enum");
 }

llvm/lib/Target/AArch64/AArch64PointerAuth.h

@@ -98,14 +98,10 @@ enum class AuthCheckMethod {
 /// using an I-key or D-key and which register can be used as temporary.
 /// If an explicit BRK instruction is used to generate an exception, BrkImm
 /// specifies its immediate operand.
-///
-/// \returns The machine basic block containing the code that is executed
-///          after the check succeeds.
-MachineBasicBlock &checkAuthenticatedRegister(MachineBasicBlock::iterator MBBI,
-                                              AuthCheckMethod Method,
-                                              Register AuthenticatedReg,
-                                              Register TmpReg, bool UseIKey,
-                                              unsigned BrkImm);
+void checkAuthenticatedRegister(MachineBasicBlock::iterator MBBI,
+                                AuthCheckMethod Method,
+                                Register AuthenticatedReg, Register TmpReg,
+                                bool UseIKey, unsigned BrkImm);
 
 /// Returns the number of bytes added by checkAuthenticatedRegister.
 unsigned getCheckerSizeInBytes(AuthCheckMethod Method);

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()