[CodeGen] Use non-preemptible symbol for BlockAddress global constants

Just as with landing pads, which are basically a special form of
BlockAddress, we should not use preemptible symbols here.

Author: jrtc27

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1107,6 +1107,35 @@ void AsmPrinter::emitFunctionHeader() {
 
 static bool needFuncLabels(const MachineFunction &MF, const AsmPrinter &Asm);
 
+/// Returns true if function begin label should be emitted due to a
+/// BlockAddress used as an initializer for a GlobalVariable.
+static bool needFuncLabelForBlockAddress(const MachineFunction &MF) {
+  SmallVector<const Value *, 4> Users;
+
+  for (const auto &MBB : MF) {
+    if (!MBB.isIRBlockAddressTaken())
+      continue;
+
+    const auto *BB = MBB.getAddressTakenIRBlock();
+    for (const auto &U : BB->users()) {
+      if (isa<BlockAddress>(U))
+        Users.push_back(U);
+    }
+  }
+
+  while (!Users.empty()) {
+    const Value *V = Users.pop_back_val();
+    if (isa<GlobalVariable>(V))
+      return true;
+    if (isa<Instruction>(V))
+      continue;
+    for (const auto &U : V->users())
+      Users.push_back(U);
+  }
+
+  return false;
+}
+
 /// EmitFunctionEntryLabel - Emit the label that is the entrypoint for the
 /// function.  This can be overridden by targets as required to do custom stuff.
 void AsmPrinter::emitFunctionEntryLabel() {
@@ -1116,7 +1145,9 @@ void AsmPrinter::emitFunctionEntryLabel() {
   if (TM.getTargetTriple().isOSBinFormatELF()) {
     // For CHERI purecap exception handling, we always have to use a local
     // alias even if the function is not dso_local.
-    bool ForceLocal = MAI->isCheriPurecapABI() && needFuncLabels(*MF, *this);
+    bool ForceLocal =
+        MAI->isCheriPurecapABI() &&
+        (needFuncLabels(*MF, *this) || needFuncLabelForBlockAddress(*MF));
     MCSymbol *Sym = getSymbolPreferLocal(MF->getFunction(), ForceLocal);
     if (Sym != CurrentFnSym) {
       CurrentFnBeginLocal = Sym;
@@ -4163,7 +4194,8 @@ static void emitGlobalConstantCHERICap(const DataLayout &DL, const Constant *CV,
   } else if (const MCSymbolRefExpr *SRE = dyn_cast<MCSymbolRefExpr>(Expr)) {
     if (auto BA = dyn_cast<BlockAddress>(CV)) {
       // For block addresses we emit `.chericap FN+(.LtmpN - FN)`
-      auto FnStart = AP.getSymbol(BA->getFunction());
+      // NB: Must use a non-preemptible symbol
+      auto FnStart = AP.getSymbolPreferLocal(*BA->getFunction(), true);
       const MCExpr *Start = MCSymbolRefExpr::create(FnStart, Ctx);
       const MCExpr *DiffToStart = MCBinaryExpr::createSub(SRE, Start, Ctx);
       const MCExpr *CapExpr = MCBinaryExpr::createAdd(Start, DiffToStart, Ctx);

llvm/test/CodeGen/Mips/cheri/cheri-blockaddress.ll

@@ -62,12 +62,12 @@ indirectgoto:                                     ; preds = %entry
 ; ASM: .end	addrof_label_in_local
 
 ; ASM-LABEL: addrof_label_in_static.b:
-; ASM-NEXT: 	.chericap	addrof_label_in_static+(.Ltmp0-addrof_label_in_static)
+; ASM-NEXT: 	.chericap	.Laddrof_label_in_static$local+(.Ltmp0-.Laddrof_label_in_static$local)
 ; ASM-NEXT: 	.size	addrof_label_in_static.b, [[#CAP_SIZE]]
 
 ; The .o file should contain a relocation against the function with a constant addend (0x1c)
 ; OBJ:       Section (8) .rela.data {
-; OBJ-NEXT:    0x0 R_MIPS_CHERI_CAPABILITY/R_MIPS_NONE/R_MIPS_NONE addrof_label_in_static 0x1C
+; OBJ-NEXT:    0x0 R_MIPS_CHERI_CAPABILITY/R_MIPS_NONE/R_MIPS_NONE .Laddrof_label_in_static$local 0x1C
 ; OBJ-NEXT:  }