Only fetch HIR for naked functions that have the attribute.

1 files changed, 41 insertions, 53 deletions

diff --git a/compiler/rustc_passes/src/naked_functions.rs b/compiler/rustc_passes/src/naked_functions.rs
index 20765abf392..296b6ed5d99 100644
--- a/compiler/rustc_passes/src/naked_functions.rs
+++ b/compiler/rustc_passes/src/naked_functions.rs
@@ -1,11 +1,12 @@
 //! Checks validity of naked functions.
 
-use rustc_ast::{Attribute, InlineAsmOptions};
+use rustc_ast::InlineAsmOptions;
 use rustc_errors::{struct_span_err, Applicability};
 use rustc_hir as hir;
+use rustc_hir::def::DefKind;
 use rustc_hir::def_id::LocalDefId;
-use rustc_hir::intravisit::{FnKind, Visitor};
-use rustc_hir::{ExprKind, HirId, InlineAsmOperand, StmtKind};
+use rustc_hir::intravisit::Visitor;
+use rustc_hir::{ExprKind, InlineAsmOperand, StmtKind};
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::lint::builtin::UNDEFINED_NAKED_FUNCTION_ABI;
@@ -13,71 +14,58 @@ use rustc_span::symbol::sym;
 use rustc_span::Span;
 use rustc_target::spec::abi::Abi;
 
-fn check_mod_naked_functions(tcx: TyCtxt<'_>, module_def_id: LocalDefId) {
-    tcx.hir().visit_item_likes_in_module(module_def_id, &mut CheckNakedFunctions { tcx });
-}
-
 pub(crate) fn provide(providers: &mut Providers) {
     *providers = Providers { check_mod_naked_functions, ..*providers };
 }
 
-struct CheckNakedFunctions<'tcx> {
-    tcx: TyCtxt<'tcx>,
-}
-
-impl<'tcx> Visitor<'tcx> for CheckNakedFunctions<'tcx> {
-    fn visit_fn(
-        &mut self,
-        fk: FnKind<'_>,
-        _fd: &'tcx hir::FnDecl<'tcx>,
-        body_id: hir::BodyId,
-        span: Span,
-        hir_id: HirId,
-    ) {
-        let ident_span;
-        let fn_header;
-
-        match fk {
-            FnKind::Closure => {
-                // Closures with a naked attribute are rejected during attribute
-                // check. Don't validate them any further.
-                return;
-            }
-            FnKind::ItemFn(ident, _, ref header, ..) => {
-                ident_span = ident.span;
-                fn_header = header;
-            }
-
-            FnKind::Method(ident, ref sig, ..) => {
-                ident_span = ident.span;
-                fn_header = &sig.header;
-            }
+fn check_mod_naked_functions(tcx: TyCtxt<'_>, module_def_id: LocalDefId) {
+    let items = tcx.hir_module_items(module_def_id);
+    for def_id in items.definitions() {
+        if !matches!(tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
+            continue;
         }
 
-        let attrs = self.tcx.hir().attrs(hir_id);
-        let naked = attrs.iter().any(|attr| attr.has_name(sym::naked));
-        if naked {
-            let body = self.tcx.hir().body(body_id);
-            check_abi(self.tcx, hir_id, fn_header.abi, ident_span);
-            check_no_patterns(self.tcx, body.params);
-            check_no_parameters_use(self.tcx, body);
-            check_asm(self.tcx, body, span);
-            check_inline(self.tcx, attrs);
+        let naked = tcx.has_attr(def_id.to_def_id(), sym::naked);
+        if !naked {
+            continue;
         }
+
+        let (fn_header, body_id) = match tcx.hir().get_by_def_id(def_id) {
+            hir::Node::Item(hir::Item { kind: hir::ItemKind::Fn(sig, _, body_id), .. })
+            | hir::Node::TraitItem(hir::TraitItem {
+                kind: hir::TraitItemKind::Fn(sig, hir::TraitFn::Provided(body_id)),
+                ..
+            })
+            | hir::Node::ImplItem(hir::ImplItem {
+                kind: hir::ImplItemKind::Fn(sig, body_id),
+                ..
+            }) => (sig.header, *body_id),
+            _ => continue,
+        };
+
+        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);
+        check_inline(tcx, def_id);
     }
 }
 
 /// Check that the function isn't inlined.
-fn check_inline(tcx: TyCtxt<'_>, attrs: &[Attribute]) {
-    for attr in attrs.iter().filter(|attr| attr.has_name(sym::inline)) {
+fn check_inline(tcx: TyCtxt<'_>, def_id: LocalDefId) {
+    let attrs = tcx.get_attrs(def_id.to_def_id(), sym::inline);
+    for attr in attrs {
         tcx.sess.struct_span_err(attr.span, "naked functions cannot be inlined").emit();
     }
 }
 
 /// Checks that function uses non-Rust ABI.
-fn check_abi(tcx: TyCtxt<'_>, hir_id: HirId, abi: Abi, fn_ident_span: Span) {
+fn check_abi(tcx: TyCtxt<'_>, def_id: LocalDefId, abi: Abi) {
     if abi == Abi::Rust {
-        tcx.struct_span_lint_hir(UNDEFINED_NAKED_FUNCTION_ABI, hir_id, fn_ident_span, |lint| {
+        let hir_id = tcx.hir().local_def_id_to_hir_id(def_id);
+        let span = tcx.def_span(def_id);
+        tcx.struct_span_lint_hir(UNDEFINED_NAKED_FUNCTION_ABI, hir_id, span, |lint| {
             lint.build("Rust ABI is unsupported in naked functions").emit();
         });
     }
@@ -141,7 +129,7 @@ impl<'tcx> Visitor<'tcx> for CheckParameters<'tcx> {
 }
 
 /// Checks that function body contains a single inline assembly block.
-fn check_asm<'tcx>(tcx: TyCtxt<'tcx>, body: &'tcx hir::Body<'tcx>, fn_span: Span) {
+fn check_asm<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId, body: &'tcx hir::Body<'tcx>) {
     let mut this = CheckInlineAssembly { tcx, items: Vec::new() };
     this.visit_body(body);
     if let [(ItemKind::Asm | ItemKind::Err, _)] = this.items[..] {
@@ -149,7 +137,7 @@ fn check_asm<'tcx>(tcx: TyCtxt<'tcx>, body: &'tcx hir::Body<'tcx>, fn_span: Span
     } else {
         let mut diag = struct_span_err!(
             tcx.sess,
-            fn_span,
+            tcx.def_span(def_id),
             E0787,
             "naked functions must contain a single asm block"
         );