Rewrite #[track_caller]

12 files changed, 82 insertions, 53 deletions

diff --git a/Cargo.lock b/Cargo.lock
index a170ece0a0d..36c8677883d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3282,6 +3282,7 @@ dependencies = [
  "rustc_abi",
  "rustc_ast",
  "rustc_ast_pretty",
+ "rustc_attr_data_structures",
  "rustc_attr_parsing",
  "rustc_data_structures",
  "rustc_errors",
diff --git a/compiler/rustc_ast_lowering/Cargo.toml b/compiler/rustc_ast_lowering/Cargo.toml
index 6ac258155fe..dc571f5c367 100644
--- a/compiler/rustc_ast_lowering/Cargo.toml
+++ b/compiler/rustc_ast_lowering/Cargo.toml
@@ -11,6 +11,7 @@ doctest = false
 rustc_abi = { path = "../rustc_abi" }
 rustc_ast = { path = "../rustc_ast" }
 rustc_ast_pretty = { path = "../rustc_ast_pretty" }
+rustc_attr_data_structures = { path = "../rustc_attr_data_structures" }
 rustc_attr_parsing = { path = "../rustc_attr_parsing" }
 rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index f297bf9f4cf..c2140514e31 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -4,6 +4,7 @@ use std::sync::Arc;
 use rustc_ast::ptr::P as AstP;
 use rustc_ast::*;
 use rustc_ast_pretty::pprust::expr_to_string;
+use rustc_attr_data_structures::{AttributeKind, find_attr};
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_hir as hir;
 use rustc_hir::HirId;
@@ -831,7 +832,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     ) {
         if self.tcx.features().async_fn_track_caller()
             && let Some(attrs) = self.attrs.get(&outer_hir_id.local_id)
-            && attrs.into_iter().any(|attr| attr.has_name(sym::track_caller))
+            && find_attr!(*attrs, AttributeKind::TrackCaller(_))
         {
             let unstable_span = self.mark_span_with_reason(
                 DesugaringKind::Async,
diff --git a/compiler/rustc_attr_data_structures/src/attributes.rs b/compiler/rustc_attr_data_structures/src/attributes.rs
index ae3a1a41387..008fc44f3cd 100644
--- a/compiler/rustc_attr_data_structures/src/attributes.rs
+++ b/compiler/rustc_attr_data_structures/src/attributes.rs
@@ -268,5 +268,8 @@ pub enum AttributeKind {
         /// Span of the attribute.
         span: Span,
     },
+
+    /// Represents `#[track_caller]`
+    TrackCaller(Span),
     // tidy-alphabetical-end
 }
diff --git a/compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs b/compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs
index 0fa69c40154..eadf8657a0f 100644
--- a/compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs
+++ b/compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs
@@ -1,7 +1,7 @@
 use rustc_attr_data_structures::{AttributeKind, OptimizeAttr};
 use rustc_feature::{AttributeTemplate, template};
 use rustc_session::parse::feature_err;
-use rustc_span::{Span, sym};
+use rustc_span::{Span, Symbol, sym};
 
 use super::{AcceptMapping, AttributeOrder, AttributeParser, OnDuplicate, SingleAttributeParser};
 use crate::context::{AcceptContext, FinalizeContext, Stage};
@@ -11,7 +11,7 @@ use crate::session_diagnostics::NakedFunctionIncompatibleAttribute;
 pub(crate) struct OptimizeParser;
 
 impl<S: Stage> SingleAttributeParser<S> for OptimizeParser {
-    const PATH: &[rustc_span::Symbol] = &[sym::optimize];
+    const PATH: &[Symbol] = &[sym::optimize];
     const ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepLast;
     const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::WarnButFutureError;
     const TEMPLATE: AttributeTemplate = template!(List: "size|speed|none");
@@ -44,7 +44,7 @@ impl<S: Stage> SingleAttributeParser<S> for OptimizeParser {
 pub(crate) struct ColdParser;
 
 impl<S: Stage> SingleAttributeParser<S> for ColdParser {
-    const PATH: &[rustc_span::Symbol] = &[sym::cold];
+    const PATH: &[Symbol] = &[sym::cold];
     const ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepLast;
     const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Warn;
     const TEMPLATE: AttributeTemplate = template!(Word);
@@ -166,6 +166,24 @@ impl<S: Stage> AttributeParser<S> for NakedParser {
     }
 }
 
+pub(crate) struct TrackCallerParser;
+
+impl<S: Stage> SingleAttributeParser<S> for TrackCallerParser {
+    const PATH: &[Symbol] = &[sym::track_caller];
+    const ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepLast;
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Warn;
+    const TEMPLATE: AttributeTemplate = template!(Word);
+
+    fn convert(cx: &mut AcceptContext<'_, '_, S>, args: &ArgParser<'_>) -> Option<AttributeKind> {
+        if let Err(span) = args.no_args() {
+            cx.expected_no_args(span);
+            return None;
+        }
+
+        Some(AttributeKind::TrackCaller(cx.attr_span))
+    }
+}
+
 pub(crate) struct NoMangleParser;
 
 impl<S: Stage> SingleAttributeParser<S> for NoMangleParser {
diff --git a/compiler/rustc_attr_parsing/src/context.rs b/compiler/rustc_attr_parsing/src/context.rs
index 43b2d07dbe1..22caf8553af 100644
--- a/compiler/rustc_attr_parsing/src/context.rs
+++ b/compiler/rustc_attr_parsing/src/context.rs
@@ -15,7 +15,9 @@ use rustc_session::Session;
 use rustc_span::{DUMMY_SP, ErrorGuaranteed, Span, Symbol, sym};
 
 use crate::attributes::allow_unstable::{AllowConstFnUnstableParser, AllowInternalUnstableParser};
-use crate::attributes::codegen_attrs::{ColdParser, NakedParser, NoMangleParser, OptimizeParser};
+use crate::attributes::codegen_attrs::{
+    ColdParser, NakedParser, NoMangleParser, OptimizeParser, TrackCallerParser,
+};
 use crate::attributes::confusables::ConfusablesParser;
 use crate::attributes::deprecation::DeprecationParser;
 use crate::attributes::inline::{InlineParser, RustcForceInlineParser};
@@ -121,6 +123,7 @@ attribute_parsers!(
         Single<PubTransparentParser>,
         Single<RustcForceInlineParser>,
         Single<SkipDuringMethodDispatchParser>,
+        Single<TrackCallerParser>,
         Single<TransparencyParser>,
         // tidy-alphabetical-end
     ];
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
index f769b393528..7bd27eb3ef1 100644
--- a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
+++ b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -95,17 +95,15 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
         // In these cases, we bail from performing further checks that are only meaningful for
         // functions (such as calling `fn_sig`, which ICEs if given a non-function). We also
         // report a delayed bug, just in case `check_attr` isn't doing its job.
-        let fn_sig = || {
+        let fn_sig = |attr_span| {
             use DefKind::*;
 
             let def_kind = tcx.def_kind(did);
             if let Fn | AssocFn | Variant | Ctor(..) = def_kind {
                 Some(tcx.fn_sig(did))
             } else {
-                tcx.dcx().span_delayed_bug(
-                    attr.span(),
-                    "this attribute can only be applied to functions",
-                );
+                tcx.dcx()
+                    .span_delayed_bug(attr_span, "this attribute can only be applied to functions");
                 None
             }
         };
@@ -142,6 +140,29 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
                         });
                     }
                 }
+                AttributeKind::TrackCaller(attr_span) => {
+                    let is_closure = tcx.is_closure_like(did.to_def_id());
+
+                    if !is_closure
+                        && let Some(fn_sig) = fn_sig(*attr_span)
+                        && fn_sig.skip_binder().abi() != ExternAbi::Rust
+                    {
+                        tcx.dcx().emit_err(errors::RequiresRustAbi { span: *attr_span });
+                    }
+                    if is_closure
+                        && !tcx.features().closure_track_caller()
+                        && !attr_span.allows_unstable(sym::closure_track_caller)
+                    {
+                        feature_err(
+                            &tcx.sess,
+                            sym::closure_track_caller,
+                            *attr_span,
+                            "`#[track_caller]` on closures is currently unstable",
+                        )
+                        .emit();
+                    }
+                    codegen_fn_attrs.flags |= CodegenFnAttrFlags::TRACK_CALLER
+                }
                 _ => {}
             }
         }
@@ -202,29 +223,6 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
                 }
             }
             sym::thread_local => codegen_fn_attrs.flags |= CodegenFnAttrFlags::THREAD_LOCAL,
-            sym::track_caller => {
-                let is_closure = tcx.is_closure_like(did.to_def_id());
-
-                if !is_closure
-                    && let Some(fn_sig) = fn_sig()
-                    && fn_sig.skip_binder().abi() != ExternAbi::Rust
-                {
-                    tcx.dcx().emit_err(errors::RequiresRustAbi { span: attr.span() });
-                }
-                if is_closure
-                    && !tcx.features().closure_track_caller()
-                    && !attr.span().allows_unstable(sym::closure_track_caller)
-                {
-                    feature_err(
-                        &tcx.sess,
-                        sym::closure_track_caller,
-                        attr.span(),
-                        "`#[track_caller]` on closures is currently unstable",
-                    )
-                    .emit();
-                }
-                codegen_fn_attrs.flags |= CodegenFnAttrFlags::TRACK_CALLER
-            }
             sym::export_name => {
                 if let Some(s) = attr.value_str() {
                     if s.as_str().contains('\0') {
diff --git a/compiler/rustc_hir_analysis/src/check/entry.rs b/compiler/rustc_hir_analysis/src/check/entry.rs
index 3bad36da999..b556683e80a 100644
--- a/compiler/rustc_hir_analysis/src/check/entry.rs
+++ b/compiler/rustc_hir_analysis/src/check/entry.rs
@@ -1,14 +1,15 @@
 use std::ops::Not;
 
 use rustc_abi::ExternAbi;
+use rustc_attr_data_structures::{AttributeKind, find_attr};
 use rustc_hir as hir;
 use rustc_hir::Node;
 use rustc_infer::infer::TyCtxtInferExt;
 use rustc_middle::span_bug;
 use rustc_middle::ty::{self, TyCtxt, TypingMode};
 use rustc_session::config::EntryFnType;
+use rustc_span::Span;
 use rustc_span::def_id::{CRATE_DEF_ID, DefId, LocalDefId};
-use rustc_span::{Span, sym};
 use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
 use rustc_trait_selection::traits::{self, ObligationCause, ObligationCauseCode};
 
@@ -98,8 +99,10 @@ fn check_main_fn_ty(tcx: TyCtxt<'_>, main_def_id: DefId) {
         error = true;
     }
 
-    for attr in tcx.get_attrs(main_def_id, sym::track_caller) {
-        tcx.dcx().emit_err(errors::TrackCallerOnMain { span: attr.span(), annotated: main_span });
+    if let Some(attr_span) =
+        find_attr!(tcx.get_all_attrs(main_def_id), AttributeKind::TrackCaller(span) => *span)
+    {
+        tcx.dcx().emit_err(errors::TrackCallerOnMain { span: attr_span, annotated: main_span });
         error = true;
     }
 
diff --git a/compiler/rustc_lint/src/builtin.rs b/compiler/rustc_lint/src/builtin.rs
index ac405277c4e..5bbe69c8d65 100644
--- a/compiler/rustc_lint/src/builtin.rs
+++ b/compiler/rustc_lint/src/builtin.rs
@@ -1184,11 +1184,11 @@ impl<'tcx> LateLintPass<'tcx> for UngatedAsyncFnTrackCaller {
         if fn_kind.asyncness().is_async()
             && !cx.tcx.features().async_fn_track_caller()
             // Now, check if the function has the `#[track_caller]` attribute
-            && let Some(attr) = cx.tcx.get_attr(def_id, sym::track_caller)
+            && let Some(attr_span) = find_attr!(cx.tcx.get_all_attrs(def_id), AttributeKind::TrackCaller(span) => *span)
         {
             cx.emit_span_lint(
                 UNGATED_ASYNC_FN_TRACK_CALLER,
-                attr.span(),
+                attr_span,
                 BuiltinUngatedAsyncFnTrackCaller { label: span, session: &cx.tcx.sess },
             );
         }
diff --git a/compiler/rustc_parse/src/validate_attr.rs b/compiler/rustc_parse/src/validate_attr.rs
index 47e29923ee5..68d78af5943 100644
--- a/compiler/rustc_parse/src/validate_attr.rs
+++ b/compiler/rustc_parse/src/validate_attr.rs
@@ -301,6 +301,7 @@ fn emit_malformed_attribute(
             | sym::naked
             | sym::no_mangle
             | sym::must_use
+            | sym::track_caller
     ) {
         return;
     }
diff --git a/compiler/rustc_passes/src/check_attr.rs b/compiler/rustc_passes/src/check_attr.rs
index 25b84179443..fc190b92816 100644
--- a/compiler/rustc_passes/src/check_attr.rs
+++ b/compiler/rustc_passes/src/check_attr.rs
@@ -169,6 +169,9 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                 Attribute::Parsed(AttributeKind::Naked(attr_span)) => {
                     self.check_naked(hir_id, *attr_span, span, target)
                 }
+                Attribute::Parsed(AttributeKind::TrackCaller(attr_span)) => {
+                    self.check_track_caller(hir_id, *attr_span, attrs, span, target)
+                }
                 Attribute::Parsed(
                     AttributeKind::BodyStability { .. }
                     | AttributeKind::ConstStabilityIndirect
@@ -205,9 +208,6 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                             self.check_target_feature(hir_id, attr, span, target, attrs)
                         }
                         [sym::thread_local, ..] => self.check_thread_local(attr, span, target),
-                        [sym::track_caller, ..] => {
-                            self.check_track_caller(hir_id, attr.span(), attrs, span, target)
-                        }
                         [sym::doc, ..] => self.check_doc_attrs(
                             attr,
                             attr_item.style,
diff --git a/tests/ui/lint/unused/unused-attr-duplicate.stderr b/tests/ui/lint/unused/unused-attr-duplicate.stderr
index 4abf7dd134e..df6cada6b06 100644
--- a/tests/ui/lint/unused/unused-attr-duplicate.stderr
+++ b/tests/ui/lint/unused/unused-attr-duplicate.stderr
@@ -90,18 +90,6 @@ LL | #[automatically_derived]
    | ^^^^^^^^^^^^^^^^^^^^^^^^
 
 error: unused attribute
-  --> $DIR/unused-attr-duplicate.rs:79:1
-   |
-LL | #[track_caller]
-   | ^^^^^^^^^^^^^^^ help: remove this attribute
-   |
-note: attribute also specified here
-  --> $DIR/unused-attr-duplicate.rs:78:1
-   |
-LL | #[track_caller]
-   | ^^^^^^^^^^^^^^^
-
-error: unused attribute
   --> $DIR/unused-attr-duplicate.rs:92:1
    |
 LL | #[export_name = "exported_symbol_name"]
@@ -278,6 +266,18 @@ LL | #[cold]
    | ^^^^^^^
 
 error: unused attribute
+  --> $DIR/unused-attr-duplicate.rs:79:1
+   |
+LL | #[track_caller]
+   | ^^^^^^^^^^^^^^^ help: remove this attribute
+   |
+note: attribute also specified here
+  --> $DIR/unused-attr-duplicate.rs:78:1
+   |
+LL | #[track_caller]
+   | ^^^^^^^^^^^^^^^
+
+error: unused attribute
   --> $DIR/unused-attr-duplicate.rs:98:1
    |
 LL | #[no_mangle]