add `codegen_instance_attrs` query

and use it for naked functions

7 files changed, 34 insertions, 23 deletions

diff --git a/compiler/rustc_codegen_cranelift/src/driver/aot.rs b/compiler/rustc_codegen_cranelift/src/driver/aot.rs
index 727f2760c0f..8ec3599b63d 100644
--- a/compiler/rustc_codegen_cranelift/src/driver/aot.rs
+++ b/compiler/rustc_codegen_cranelift/src/driver/aot.rs
@@ -530,12 +530,8 @@ fn codegen_cgu_content(
     for (mono_item, item_data) in mono_items {
         match mono_item {
             MonoItem::Fn(instance) => {
-                // Other `InstanceKind`s (e.g. `ReifyShim` generated by indirect calls) should be
-                // codegened like a normal function.
-                let is_item_instance = matches!(instance.def, InstanceKind::Item(_));
-
-                let flags = tcx.codegen_fn_attrs(instance.def_id()).flags;
-                if is_item_instance && flags.contains(CodegenFnAttrFlags::NAKED) {
+                let flags = tcx.codegen_instance_attrs(instance.def).flags;
+                if flags.contains(CodegenFnAttrFlags::NAKED) {
                     rustc_codegen_ssa::mir::naked_asm::codegen_naked_asm(
                         &mut GlobalAsmContext { tcx, global_asm: &mut cx.global_asm },
                         instance,
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
index 85d01d4f938..a1dd64736fc 100644
--- a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
+++ b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -114,6 +114,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
                     codegen_fn_attrs.export_name = Some(*name);
                 }
                 AttributeKind::Naked(_) => codegen_fn_attrs.flags |= CodegenFnAttrFlags::NAKED,
+
                 AttributeKind::Align { align, .. } => codegen_fn_attrs.alignment = Some(*align),
                 AttributeKind::LinkName { name, .. } => codegen_fn_attrs.link_name = Some(*name),
                 AttributeKind::LinkSection { name, .. } => {
diff --git a/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs b/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
index 025f5fb54f4..b8f635ab781 100644
--- a/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
@@ -356,7 +356,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             LocalRef::Operand(operand) => {
                 // Don't spill operands onto the stack in naked functions.
                 // See: https://github.com/rust-lang/rust/issues/42779
-                let attrs = bx.tcx().codegen_fn_attrs(self.instance.def_id());
+                let attrs = bx.tcx().codegen_instance_attrs(self.instance.def);
                 if attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
                     return;
                 }
diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index fa69820d5d2..50d0f910744 100644
--- a/compiler/rustc_codegen_ssa/src/mir/mod.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/mod.rs
@@ -390,9 +390,8 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     let mut num_untupled = None;
 
-    let codegen_fn_attrs = bx.tcx().codegen_fn_attrs(fx.instance.def_id());
-    let naked = codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED);
-    if naked {
+    let codegen_fn_attrs = bx.tcx().codegen_instance_attrs(fx.instance.def);
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NAKED) {
         return vec![];
     }
 
diff --git a/compiler/rustc_codegen_ssa/src/mono_item.rs b/compiler/rustc_codegen_ssa/src/mono_item.rs
index fa0ef04b2e6..2c3a6de750e 100644
--- a/compiler/rustc_codegen_ssa/src/mono_item.rs
+++ b/compiler/rustc_codegen_ssa/src/mono_item.rs
@@ -1,6 +1,5 @@
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::mono::{Linkage, MonoItem, MonoItemData, Visibility};
-use rustc_middle::ty::InstanceKind;
 use rustc_middle::ty::layout::HasTyCtxt;
 use tracing::debug;
 
@@ -42,12 +41,8 @@ impl<'a, 'tcx: 'a> MonoItemExt<'a, 'tcx> for MonoItem<'tcx> {
                 base::codegen_global_asm(cx, item_id);
             }
             MonoItem::Fn(instance) => {
-                // Other `InstanceKind`s (e.g. `ReifyShim` generated by indirect calls) should be
-                // codegened like a normal function.
-                let is_item_instance = matches!(instance.def, InstanceKind::Item(_));
-
-                let flags = cx.tcx().codegen_fn_attrs(instance.def_id()).flags;
-                if is_item_instance && flags.contains(CodegenFnAttrFlags::NAKED) {
+                let flags = cx.tcx().codegen_instance_attrs(instance.def).flags;
+                if flags.contains(CodegenFnAttrFlags::NAKED) {
                     naked_asm::codegen_naked_asm::<Bx::CodegenCx>(cx, instance, item_data);
                 } else {
                     base::codegen_instance::<Bx>(cx, instance);
diff --git a/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs b/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
index 2f16d385efb..92670e7e018 100644
--- a/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
+++ b/compiler/rustc_middle/src/middle/codegen_fn_attrs.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use rustc_abi::Align;
 use rustc_ast::expand::autodiff_attrs::AutoDiffAttrs;
 use rustc_attr_data_structures::{InlineAttr, InstructionSetAttr, OptimizeAttr};
@@ -6,6 +8,24 @@ use rustc_span::Symbol;
 use rustc_target::spec::SanitizerSet;
 
 use crate::mir::mono::Linkage;
+use crate::ty::{InstanceKind, TyCtxt};
+
+impl<'tcx> TyCtxt<'tcx> {
+    pub fn codegen_instance_attrs(
+        self,
+        instance_kind: InstanceKind<'_>,
+    ) -> Cow<'tcx, CodegenFnAttrs> {
+        let mut attrs = Cow::Borrowed(self.codegen_fn_attrs(instance_kind.def_id()));
+
+        // Drop the `#[naked]` attribute on non-item `InstanceKind`s, like the shims that
+        // are generated for indirect function calls.
+        if !matches!(instance_kind, InstanceKind::Item(_)) {
+            attrs.to_mut().flags.remove(CodegenFnAttrFlags::NAKED);
+        }
+
+        attrs
+    }
+}
 
 #[derive(Clone, TyEncodable, TyDecodable, HashStable, Debug)]
 pub struct CodegenFnAttrs {
diff --git a/tests/codegen/sanitizer/kcfi/naked-function.rs b/tests/codegen/sanitizer/kcfi/naked-function.rs
index 92431fc7d63..2c8cdc919b8 100644
--- a/tests/codegen/sanitizer/kcfi/naked-function.rs
+++ b/tests/codegen/sanitizer/kcfi/naked-function.rs
@@ -25,6 +25,11 @@ trait MyTrait {
 }
 impl MyTrait for Thing {}
 
+// the shim calls the real function
+// CHECK-LABEL: define
+// CHECK-SAME: my_naked_function
+// CHECK-SAME: reify.shim.fnptr
+
 // CHECK-LABEL: main
 #[unsafe(no_mangle)]
 pub fn main() {
@@ -32,16 +37,11 @@ pub fn main() {
     const F: extern "C" fn() = Thing::my_naked_function;
 
     // main calls the shim function
-    // CHECK: call
+    // CHECK: call void
     // CHECK-SAME: my_naked_function
     // CHECK-SAME: reify.shim.fnptr
     (F)();
 }
 
-// the shim calls the real function
-// CHECK: define
-// CHECK-SAME: my_naked_function
-// CHECK-SAME: reify.shim.fnptr
-
 // CHECK: declare !kcfi_type
 // CHECK-SAME: my_naked_function