Auto merge of #130195 - folkertdev:naked-asm-outside-naked-fn, r=Amanieu

disallow `naked_asm!` outside of `#[naked]` functions

tracking issue: #90957
parent PR: #128651

I split this out from the parent PR because it's self-contained and because the analysis has to search through all functions and there might be performance regressions.

r? `@Amanieu`

Author: bors

Diff for: compiler/rustc_ast/src/ast.rs

@@ -2418,11 +2418,22 @@ impl InlineAsmOperand {
     }
 }
 
+#[derive(Clone, Copy, Encodable, Decodable, Debug, HashStable_Generic)]
+pub enum AsmMacro {
+    /// The `asm!` macro
+    Asm,
+    /// The `global_asm!` macro
+    GlobalAsm,
+    /// The `naked_asm!` macro
+    NakedAsm,
+}
+
 /// Inline assembly.
 ///
 /// E.g., `asm!("NOP");`.
 #[derive(Clone, Encodable, Decodable, Debug)]
 pub struct InlineAsm {
+    pub asm_macro: AsmMacro,
     pub template: Vec<InlineAsmTemplatePiece>,
     pub template_strs: Box<[(Symbol, Option<Symbol>, Span)]>,
     pub operands: Vec<(InlineAsmOperand, Span)>,

Diff for: compiler/rustc_ast/src/mut_visit.rs

@@ -1388,6 +1388,7 @@ fn walk_anon_const<T: MutVisitor>(vis: &mut T, AnonConst { id, value }: &mut Ano
 fn walk_inline_asm<T: MutVisitor>(vis: &mut T, asm: &mut InlineAsm) {
     // FIXME: Visit spans inside all this currently ignored stuff.
     let InlineAsm {
+        asm_macro: _,
         template: _,
         template_strs: _,
         operands,

Diff for: compiler/rustc_ast/src/visit.rs

@@ -976,6 +976,7 @@ pub fn walk_anon_const<'a, V: Visitor<'a>>(visitor: &mut V, constant: &'a AnonCo
 
 pub fn walk_inline_asm<'a, V: Visitor<'a>>(visitor: &mut V, asm: &'a InlineAsm) -> V::Result {
     let InlineAsm {
+        asm_macro: _,
         template: _,
         template_strs: _,
         operands,

Diff for: compiler/rustc_ast_lowering/src/asm.rs

@@ -474,8 +474,14 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         );
         let line_spans =
             self.arena.alloc_from_iter(asm.line_spans.iter().map(|span| self.lower_span(*span)));
-        let hir_asm =
-            hir::InlineAsm { template, template_strs, operands, options: asm.options, line_spans };
+        let hir_asm = hir::InlineAsm {
+            asm_macro: asm.asm_macro,
+            template,
+            template_strs,
+            operands,
+            options: asm.options,
+            line_spans,
+        };
         self.arena.alloc(hir_asm)
     }
 }

Diff for: compiler/rustc_builtin_macros/src/asm.rs

@@ -3,6 +3,7 @@ use lint::BuiltinLintDiag;
 use rustc_ast::ptr::P;
 use rustc_ast::token::{self, Delimiter};
 use rustc_ast::tokenstream::TokenStream;
+use rustc_ast::AsmMacro;
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_errors::PResult;
 use rustc_expand::base::*;
@@ -484,6 +485,7 @@ fn parse_reg<'a>(
 
 fn expand_preparsed_asm(
     ecx: &mut ExtCtxt<'_>,
+    asm_macro: ast::AsmMacro,
     args: AsmArgs,
 ) -> ExpandResult<Result<ast::InlineAsm, ErrorGuaranteed>, ()> {
     let mut template = vec![];
@@ -774,6 +776,7 @@ fn expand_preparsed_asm(
     }
 
     ExpandResult::Ready(Ok(ast::InlineAsm {
+        asm_macro,
         template,
         template_strs: template_strs.into_boxed_slice(),
         operands: args.operands,
@@ -790,7 +793,7 @@ pub(super) fn expand_asm<'cx>(
 ) -> MacroExpanderResult<'cx> {
     ExpandResult::Ready(match parse_args(ecx, sp, tts, false) {
         Ok(args) => {
-            let ExpandResult::Ready(mac) = expand_preparsed_asm(ecx, args) else {
+            let ExpandResult::Ready(mac) = expand_preparsed_asm(ecx, AsmMacro::Asm, args) else {
                 return ExpandResult::Retry(());
             };
             let expr = match mac {
@@ -819,7 +822,8 @@ pub(super) fn expand_naked_asm<'cx>(
 ) -> MacroExpanderResult<'cx> {
     ExpandResult::Ready(match parse_args(ecx, sp, tts, false) {
         Ok(args) => {
-            let ExpandResult::Ready(mac) = expand_preparsed_asm(ecx, args) else {
+            let ExpandResult::Ready(mac) = expand_preparsed_asm(ecx, AsmMacro::NakedAsm, args)
+            else {
                 return ExpandResult::Retry(());
             };
             let expr = match mac {
@@ -857,7 +861,8 @@ pub(super) fn expand_global_asm<'cx>(
 ) -> MacroExpanderResult<'cx> {
     ExpandResult::Ready(match parse_args(ecx, sp, tts, true) {
         Ok(args) => {
-            let ExpandResult::Ready(mac) = expand_preparsed_asm(ecx, args) else {
+            let ExpandResult::Ready(mac) = expand_preparsed_asm(ecx, AsmMacro::GlobalAsm, args)
+            else {
                 return ExpandResult::Retry(());
             };
             match mac {

Diff for: compiler/rustc_hir/src/hir.rs

@@ -2927,6 +2927,7 @@ impl<'hir> InlineAsmOperand<'hir> {
 
 #[derive(Debug, Clone, Copy, HashStable_Generic)]
 pub struct InlineAsm<'hir> {
+    pub asm_macro: ast::AsmMacro,
     pub template: &'hir [InlineAsmTemplatePiece],
     pub template_strs: &'hir [(Symbol, Option<Symbol>, Span)],
     pub operands: &'hir [(InlineAsmOperand<'hir>, Span)],

Diff for: compiler/rustc_passes/messages.ftl

@@ -484,6 +484,9 @@ passes_must_not_suspend =
 passes_must_use_no_effect =
     `#[must_use]` has no effect when applied to {$article} {$target}
 
+passes_naked_asm_outside_naked_fn =
+    the `naked_asm!` macro can only be used in functions marked with `#[naked]`
+
 passes_naked_functions_asm_block =
     naked functions must contain a single asm block
     .label_multiple_asm = multiple asm blocks are unsupported in naked functions

Diff for: compiler/rustc_passes/src/errors.rs

@@ -1221,6 +1221,13 @@ pub(crate) struct NakedFunctionIncompatibleAttribute {
     pub attr: Symbol,
 }
 
+#[derive(Diagnostic)]
+#[diag(passes_naked_asm_outside_naked_fn)]
+pub(crate) struct NakedAsmOutsideNakedFn {
+    #[primary_span]
+    pub span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(passes_attr_only_in_functions)]
 pub(crate) struct AttrOnlyInFunctions {

Diff for: compiler/rustc_passes/src/naked_functions.rs

@@ -6,6 +6,7 @@ use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{LocalDefId, LocalModDefId};
 use rustc_hir::intravisit::Visitor;
 use rustc_hir::{ExprKind, HirIdSet, InlineAsmOperand, StmtKind};
+use rustc_middle::hir::nested_filter::OnlyBodies;
 use rustc_middle::query::Providers;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::lint::builtin::UNDEFINED_NAKED_FUNCTION_ABI;
@@ -14,8 +15,9 @@ use rustc_span::Span;
 use rustc_target::spec::abi::Abi;
 
 use crate::errors::{
-    NakedFunctionsAsmBlock, NakedFunctionsAsmOptions, NakedFunctionsMustUseNoreturn,
-    NakedFunctionsOperands, NoPatterns, ParamsNotAllowed, UndefinedNakedFunctionAbi,
+    NakedAsmOutsideNakedFn, NakedFunctionsAsmBlock, NakedFunctionsAsmOptions,
+    NakedFunctionsMustUseNoreturn, NakedFunctionsOperands, NoPatterns, ParamsNotAllowed,
+    UndefinedNakedFunctionAbi,
 };
 
 pub(crate) fn provide(providers: &mut Providers) {
@@ -29,11 +31,6 @@ fn check_mod_naked_functions(tcx: TyCtxt<'_>, module_def_id: LocalModDefId) {
             continue;
         }
 
-        let naked = tcx.has_attr(def_id, sym::naked);
-        if !naked {
-            continue;
-        }
-
         let (fn_header, body_id) = match tcx.hir_node_by_def_id(def_id) {
             hir::Node::Item(hir::Item { kind: hir::ItemKind::Fn(sig, _, body_id), .. })
             | hir::Node::TraitItem(hir::TraitItem {
@@ -48,10 +45,17 @@ fn check_mod_naked_functions(tcx: TyCtxt<'_>, module_def_id: LocalModDefId) {
         };
 
         let body = tcx.hir().body(body_id);
-        check_abi(tcx, def_id, fn_header.abi);
-        check_no_patterns(tcx, body.params);
-        check_no_parameters_use(tcx, body);
-        check_asm(tcx, def_id, body);
+
+        if tcx.has_attr(def_id, sym::naked) {
+            check_abi(tcx, def_id, fn_header.abi);
+            check_no_patterns(tcx, body.params);
+            check_no_parameters_use(tcx, body);
+            check_asm(tcx, def_id, body);
+        } else {
+            // `naked_asm!` is not allowed outside of functions marked as `#[naked]`
+            let mut visitor = CheckNakedAsmInNakedFn { tcx };
+            visitor.visit_body(body);
+        }
     }
 }
 
@@ -276,3 +280,25 @@ impl<'tcx> Visitor<'tcx> for CheckInlineAssembly<'tcx> {
         self.check_expr(expr, expr.span);
     }
 }
+
+struct CheckNakedAsmInNakedFn<'tcx> {
+    tcx: TyCtxt<'tcx>,
+}
+
+impl<'tcx> Visitor<'tcx> for CheckNakedAsmInNakedFn<'tcx> {
+    type NestedFilter = OnlyBodies;
+
+    fn nested_visit_map(&mut self) -> Self::Map {
+        self.tcx.hir()
+    }
+
+    fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) {
+        if let ExprKind::InlineAsm(inline_asm) = expr.kind {
+            if let rustc_ast::AsmMacro::NakedAsm = inline_asm.asm_macro {
+                self.tcx.dcx().emit_err(NakedAsmOutsideNakedFn { span: expr.span });
+            }
+        }
+
+        hir::intravisit::walk_expr(self, expr);
+    }
+}

Diff for: tests/ui/asm/naked-asm-outside-naked-fn.rs

@@ -0,0 +1,35 @@
+//@ edition: 2021
+//@ needs-asm-support
+//@ ignore-nvptx64
+//@ ignore-spirv
+
+#![feature(naked_functions)]
+#![crate_type = "lib"]
+
+use std::arch::naked_asm;
+
+fn main() {
+    test1();
+}
+
+#[naked]
+extern "C" fn test1() {
+    unsafe { naked_asm!("") }
+}
+
+extern "C" fn test2() {
+    unsafe { naked_asm!("") }
+    //~^ ERROR the `naked_asm!` macro can only be used in functions marked with `#[naked]`
+}
+
+extern "C" fn test3() {
+    unsafe { (|| naked_asm!(""))() }
+    //~^ ERROR the `naked_asm!` macro can only be used in functions marked with `#[naked]`
+}
+
+fn test4() {
+    async move {
+        unsafe {  naked_asm!("") } ;
+        //~^ ERROR the `naked_asm!` macro can only be used in functions marked with `#[naked]`
+    };
+}

Diff for: tests/ui/asm/naked-asm-outside-naked-fn.stderr

@@ -0,0 +1,20 @@
+error: the `naked_asm!` macro can only be used in functions marked with `#[naked]`
+  --> $DIR/naked-asm-outside-naked-fn.rs:21:14
+   |
+LL |     unsafe { naked_asm!("") }
+   |              ^^^^^^^^^^^^^^
+
+error: the `naked_asm!` macro can only be used in functions marked with `#[naked]`
+  --> $DIR/naked-asm-outside-naked-fn.rs:26:18
+   |
+LL |     unsafe { (|| naked_asm!(""))() }
+   |                  ^^^^^^^^^^^^^^
+
+error: the `naked_asm!` macro can only be used in functions marked with `#[naked]`
+  --> $DIR/naked-asm-outside-naked-fn.rs:32:19
+   |
+LL |         unsafe {  naked_asm!("") } ;
+   |                   ^^^^^^^^^^^^^^
+
+error: aborting due to 3 previous errors
+