-fsanitize=function: fix .subsections_via_symbols

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -868,6 +868,9 @@ class AsmPrinter : public MachineFunctionPass {
   /// This method emits a comment next to header for the current function.
   virtual void emitFunctionHeaderComment();
 
+  /// This method emits prefix-like data before the current function.
+  void emitFunctionPrefix(const SmallVector<const Constant *, 1> &Prefix);
+
   /// Emit a blob of inline asm to the output streamer.
   void
   emitInlineAsm(StringRef Str, const MCSubtargetInfo &STI,

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -927,6 +927,30 @@ void AsmPrinter::emitDebugValue(const MCExpr *Value, unsigned Size) const {
 
 void AsmPrinter::emitFunctionHeaderComment() {}
 
+void AsmPrinter::emitFunctionPrefix(
+    const SmallVector<const Constant *, 1> &Prefix) {
+  const Function &F = MF->getFunction();
+  if (!MAI->hasSubsectionsViaSymbols()) {
+    for (auto &C : Prefix) {
+      emitGlobalConstant(F.getParent()->getDataLayout(), C);
+    }
+    return;
+  }
+  // Preserving prefix data on platforms which use subsections-via-symbols
+  // is a bit tricky. Here we introduce a symbol for the prefix data
+  // and use the .alt_entry attribute to mark the function's real entry point
+  // as an alternative entry point to the prefix-data symbol.
+  MCSymbol *PrefixSym = OutContext.createLinkerPrivateTempSymbol();
+  OutStreamer->emitLabel(PrefixSym);
+
+  for (auto &C : Prefix) {
+    emitGlobalConstant(F.getParent()->getDataLayout(), C);
+  }
+
+  // Emit an .alt_entry directive for the actual function symbol.
+  OutStreamer->emitSymbolAttribute(CurrentFnSym, MCSA_AltEntry);
+}
+
 /// EmitFunctionHeader - This method emits the header for the current
 /// function.
 void AsmPrinter::emitFunctionHeader() {
@@ -967,21 +991,7 @@ void AsmPrinter::emitFunctionHeader() {
 
   // Emit the prefix data.
   if (F.hasPrefixData()) {
-    if (MAI->hasSubsectionsViaSymbols()) {
-      // Preserving prefix data on platforms which use subsections-via-symbols
-      // is a bit tricky. Here we introduce a symbol for the prefix data
-      // and use the .alt_entry attribute to mark the function's real entry point
-      // as an alternative entry point to the prefix-data symbol.
-      MCSymbol *PrefixSym = OutContext.createLinkerPrivateTempSymbol();
-      OutStreamer->emitLabel(PrefixSym);
-
-      emitGlobalConstant(F.getParent()->getDataLayout(), F.getPrefixData());
-
-      // Emit an .alt_entry directive for the actual function symbol.
-      OutStreamer->emitSymbolAttribute(CurrentFnSym, MCSA_AltEntry);
-    } else {
-      emitGlobalConstant(F.getParent()->getDataLayout(), F.getPrefixData());
-    }
+    emitFunctionPrefix({F.getPrefixData()});
   }
 
   // Emit KCFI type information before patchable-function-prefix nops.
@@ -1014,8 +1024,7 @@ void AsmPrinter::emitFunctionHeader() {
 
     auto *PrologueSig = mdconst::extract<Constant>(MD->getOperand(0));
     auto *TypeHash = mdconst::extract<Constant>(MD->getOperand(1));
-    emitGlobalConstant(F.getParent()->getDataLayout(), PrologueSig);
-    emitGlobalConstant(F.getParent()->getDataLayout(), TypeHash);
+    emitFunctionPrefix({PrologueSig, TypeHash});
   }
 
   if (isVerbose()) {

llvm/test/CodeGen/AArch64/func-sanitizer.ll

@@ -1,4 +1,5 @@
 ; RUN: llc -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s
+; RUN: llc -mtriple=arm64-apple-darwin < %s | FileCheck %s --check-prefix=MACHO
 
 ; CHECK-LABEL: .type _Z3funv,@function
 ; CHECK-NEXT:    .word   3238382334  // 0xc105cafe
@@ -7,6 +8,14 @@
 ; CHECK-NEXT:  // %bb.0:
 ; CHECK-NEXT:    ret
 
+; MACHO: ltmp0
+; MACHO-NEXT:   .long 3238382334 ; 0xc105cafe
+; MACHO-NEXT:   .long 42 ; 0x2a
+; MACHO-NEXT:   .alt_entry __Z3funv
+; MACHO-NEXT:   __Z3funv:
+; MACHO-NEXT:   ; %bb.0:
+; MACHO-NEXT:   ret
+
 define dso_local void @_Z3funv() nounwind !func_sanitize !0 {
   ret void
 }

llvm/test/CodeGen/X86/func-sanitizer.ll

@@ -1,4 +1,5 @@
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+; RUN: llc -mtriple=x86_64-apple-darwin < %s | FileCheck %s --check-prefix=MACHO
 
 ; CHECK:      .type _Z3funv,@function
 ; CHECK-NEXT:   .long   3238382334  # 0xc105cafe
@@ -8,6 +9,15 @@
 ; CHECK-NEXT:   # %bb.0:
 ; CHECK-NEXT:   retq
 
+; MACHO:      ltmp0
+; MACHO-NEXT:  .long 3238382334 ## 0xc105cafe
+; MACHO-NEXT:  .long 42 ## 0x2a
+; MACHO-NEXT:  .alt_entry __Z3funv
+; MACHO-NEXT: __Z3funv:
+; MACHO-NEXT:  .cfi_startproc
+; MACHO-NEXT:  # %bb.0:
+; MACHO-NEXT:  retq
+
 define dso_local void @_Z3funv() !func_sanitize !0 {
   ret void
 }