diff --git a/compiler/rustc_public/src/compiler_interface.rs b/compiler/rustc_public/src/compiler_interface.rs index 5512371b68b3a..1f2a038765d4b 100644 --- a/compiler/rustc_public/src/compiler_interface.rs +++ b/compiler/rustc_public/src/compiler_interface.rs @@ -496,6 +496,14 @@ impl<'tcx> CompilerInterface<'tcx> { }) } + /// Create a caller location constant from a span. + pub(crate) fn span_as_caller_location(&self, span: Span) -> MirConst { + self.with_cx(|tables, cx| { + let sp = tables.spans[span]; + cx.span_as_caller_location(sp).stable(tables, cx) + }) + } + /// Create a new constant that represents the given string value. pub(crate) fn new_const_str(&self, value: &str) -> MirConst { self.with_cx(|tables, cx| cx.new_const_str(value).stable(tables, cx)) @@ -637,6 +645,14 @@ impl<'tcx> CompilerInterface<'tcx> { }) } + /// Check if this instance requires a caller location argument. + pub(crate) fn instance_requires_caller_location(&self, def: InstanceDef) -> bool { + self.with_cx(|tables, cx| { + let instance = tables.instances[def]; + cx.instance_requires_caller_location(instance) + }) + } + /// Check if this is an empty DropGlue shim. pub(crate) fn is_empty_drop_shim(&self, def: InstanceDef) -> bool { self.with_cx(|tables, cx| { diff --git a/compiler/rustc_public/src/mir/body.rs b/compiler/rustc_public/src/mir/body.rs index 20682de5f4825..def6c837a2b85 100644 --- a/compiler/rustc_public/src/mir/body.rs +++ b/compiler/rustc_public/src/mir/body.rs @@ -20,10 +20,10 @@ pub struct Body { /// The first local is the return value pointer, followed by `arg_count` /// locals for the function arguments, followed by any user-declared /// variables and temporaries. - pub(super) locals: LocalDecls, + pub(crate) locals: LocalDecls, /// The number of arguments this function takes. - pub(super) arg_count: usize, + pub(crate) arg_count: usize, /// Debug information pertaining to user variables, including captures. pub var_debug_info: Vec, @@ -31,19 +31,29 @@ pub struct Body { /// Mark an argument (which must be a tuple) as getting passed as its individual components. /// /// This is used for the "rust-call" ABI such as closures. - pub(super) spread_arg: Option, + pub(crate) spread_arg: Option, /// The span that covers the entire function body. pub span: Span, + + /// Source scope information, used by [`Body::caller_location`] for inline-aware resolution. + /// + /// Invariants: + /// - All scope indices referenced by terminators and statements must be within bounds. + /// - `inlined_parent_scope` links must not form cycles. + pub(crate) source_scopes: Vec, } pub type BasicBlockIdx = usize; impl Body { - /// Constructs a `Body`. + /// Constructs a `Body` without inlining information. + /// + /// # Warning /// - /// A constructor is required to build a `Body` from outside the crate - /// because the `arg_count` and `locals` fields are private. + /// This constructor does not accept source scope data today. + /// [`Body::caller_location`] will fall back to the terminator's span, + /// which may be incorrect when MIR inlining is involved. pub fn new( blocks: Vec, locals: LocalDecls, @@ -52,13 +62,15 @@ impl Body { spread_arg: Option, span: Span, ) -> Self { - // If locals doesn't contain enough entries, it can lead to panics in - // `ret_local`, `arg_locals`, and `inner_locals`. assert!( locals.len() > arg_count, "A Body must contain at least a local for the return value and each of the function's arguments" ); - Self { blocks, locals, arg_count, var_debug_info, spread_arg, span } + let source_scopes = vec![ + SourceScopeInfo { inlined: None, inlined_parent_scope: None }; + max_scope(&blocks) as usize + 1 + ]; + Self { blocks, locals, arg_count, var_debug_info, spread_arg, span, source_scopes } } /// Return local that holds this function's return value. @@ -119,6 +131,44 @@ impl Body { pub fn spread_arg(&self) -> Option { self.spread_arg } + + /// Resolve the caller location for a call to a `#[track_caller]` function. + /// + /// Use this when generating the implicit `&'static Location<'static>` argument + /// for a call where [`Instance::requires_caller_location`] is true. + /// + /// Pass `inherited_location` if this body belongs to a `#[track_caller]` function + /// (the implicit parameter it received). Pass `None` otherwise. + /// + /// This method accounts for MIR inlining: when inlined `#[track_caller]` functions + /// are present, the terminator's span may not be the correct location. The method + /// walks the inlined scopes to resolve the right one. + /// + /// [`Instance::requires_caller_location`]: crate::mir::mono::Instance::requires_caller_location + pub fn caller_location( + &self, + terminator: &Terminator, + inherited_location: Option, + ) -> MirConst { + let mut span = terminator.source_info.span; + let mut scope = terminator.source_info.scope; + + while let Some(scope_data) = self.source_scopes.get(scope as usize) { + if let Some((track_caller, callsite_span)) = &scope_data.inlined { + if !track_caller { + return span.as_caller_location(); + } + span = *callsite_span; + } + + match scope_data.inlined_parent_scope { + Some(parent) => scope = parent, + None => break, + } + } + + inherited_location.unwrap_or_else(|| span.as_caller_location()) + } } type LocalDecls = Vec; @@ -748,6 +798,22 @@ impl VarDebugInfo { pub type SourceScope = u32; +/// Data about a source scope, used for caller location resolution. +/// +/// Each entry corresponds to a source scope in the MIR body. Most scopes have no +/// inlined data. For scopes introduced by MIR inlining, `inlined` records whether +/// the inlined callee is `#[track_caller]` and the span of the call site. +#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize)] +pub(crate) struct SourceScopeInfo { + /// Present when this scope was introduced by inlining a function. + /// The `bool` is `true` if the inlined callee is `#[track_caller]`. + /// The `Span` is the call site where inlining occurred. + pub inlined: Option<(bool, Span)>, + /// Nearest (transitive) parent scope that was itself inlined. + /// Skips over intermediate scopes within the same inlined function body. + pub inlined_parent_scope: Option, +} + #[derive(Clone, Debug, Eq, PartialEq, Serialize)] pub struct SourceInfo { pub span: Span, @@ -1103,3 +1169,15 @@ impl ProjectionElem { Ok(deref_ty.ty) } } + +/// Return the maximum scope index referenced by any terminator or statement in `blocks`. +fn max_scope(blocks: &[BasicBlock]) -> u32 { + blocks + .iter() + .flat_map(|bb| { + std::iter::once(bb.terminator.source_info.scope) + .chain(bb.statements.iter().map(|s| s.source_info.scope)) + }) + .max() + .unwrap_or(0) +} diff --git a/compiler/rustc_public/src/mir/mono.rs b/compiler/rustc_public/src/mir/mono.rs index ab939a5535149..4abbb004a8e24 100644 --- a/compiler/rustc_public/src/mir/mono.rs +++ b/compiler/rustc_public/src/mir/mono.rs @@ -166,6 +166,18 @@ impl Instance { self.kind == InstanceKind::Shim && with(|cx| cx.is_empty_drop_shim(self.def)) } + /// Check whether this instance requires a caller location argument. + /// + /// Functions annotated with `#[track_caller]` have an implicit extra + /// `&'static core::panic::Location<'static>` argument appended to their ABI. + /// This argument is not present in the MIR body's signature. + /// + /// When this returns `true`, the instance's `fn_abi()` will have one additional + /// argument compared to the MIR body's parameter list. + pub fn requires_caller_location(&self) -> bool { + with(|cx| cx.instance_requires_caller_location(self.def)) + } + /// Try to constant evaluate the instance into a constant with the given type. /// /// This can be used to retrieve a constant that represents an intrinsic return such as diff --git a/compiler/rustc_public/src/mir/visit.rs b/compiler/rustc_public/src/mir/visit.rs index 8cb3237bc4c83..d80c769cd8169 100644 --- a/compiler/rustc_public/src/mir/visit.rs +++ b/compiler/rustc_public/src/mir/visit.rs @@ -408,7 +408,15 @@ macro_rules! super_body { }; ($self:ident, $body:ident, ) => { - let Body { blocks, locals: _, arg_count, var_debug_info, spread_arg: _, span } = $body; + let Body { + blocks, + locals: _, + arg_count, + var_debug_info, + spread_arg: _, + span, + source_scopes: _, + } = $body; for bb in blocks { $self.visit_basic_block(bb); diff --git a/compiler/rustc_public/src/ty.rs b/compiler/rustc_public/src/ty.rs index 8a739df7d9ad2..359a4ab62b185 100644 --- a/compiler/rustc_public/src/ty.rs +++ b/compiler/rustc_public/src/ty.rs @@ -287,6 +287,11 @@ impl Span { pub fn diagnostic(&self) -> String { with(|c| c.span_to_string(*self)) } + + /// Create a `&'static core::panic::Location<'static>` constant from this span. + pub(crate) fn as_caller_location(&self) -> MirConst { + with(|c| c.span_as_caller_location(*self)) + } } #[derive(Clone, Copy, Debug, Serialize)] diff --git a/compiler/rustc_public/src/unstable/convert/stable/mir.rs b/compiler/rustc_public/src/unstable/convert/stable/mir.rs index 4293d1bede421..6471b106ae757 100644 --- a/compiler/rustc_public/src/unstable/convert/stable/mir.rs +++ b/compiler/rustc_public/src/unstable/convert/stable/mir.rs @@ -7,7 +7,9 @@ use rustc_public_bridge::{Tables, bridge}; use crate::compiler_interface::BridgeTys; use crate::mir::alloc::GlobalAlloc; -use crate::mir::{ConstOperand, Statement, UserTypeProjection, VarDebugInfoFragment}; +use crate::mir::{ + ConstOperand, SourceScopeInfo, Statement, UserTypeProjection, VarDebugInfoFragment, +}; use crate::ty::{Allocation, ConstantKind, MirConst}; use crate::unstable::Stable; use crate::{Error, alloc, opaque}; @@ -20,8 +22,9 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> { tables: &mut Tables<'cx, BridgeTys>, cx: &CompilerCtxt<'cx, BridgeTys>, ) -> Self::T { - crate::mir::Body::new( - self.basic_blocks + crate::mir::Body { + blocks: self + .basic_blocks .iter() .map(|block| crate::mir::BasicBlock { terminator: block.terminator().stable(tables, cx), @@ -32,7 +35,8 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> { .collect(), }) .collect(), - self.local_decls + locals: self + .local_decls .iter() .map(|decl| crate::mir::LocalDecl { ty: decl.ty.stable(tables, cx), @@ -40,11 +44,25 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> { mutability: decl.mutability.stable(tables, cx), }) .collect(), - self.arg_count, - self.var_debug_info.iter().map(|info| info.stable(tables, cx)).collect(), - self.spread_arg.stable(tables, cx), - self.span.stable(tables, cx), - ) + arg_count: self.arg_count, + var_debug_info: self + .var_debug_info + .iter() + .map(|info| info.stable(tables, cx)) + .collect(), + spread_arg: self.spread_arg.stable(tables, cx), + span: self.span.stable(tables, cx), + source_scopes: self + .source_scopes + .iter() + .map(|scope_data| SourceScopeInfo { + inlined: scope_data.inlined.map(|(instance, span)| { + (instance.def.requires_caller_location(cx.tcx), span.stable(tables, cx)) + }), + inlined_parent_scope: scope_data.inlined_parent_scope.map(|s| s.as_u32()), + }) + .collect(), + } } } diff --git a/compiler/rustc_public_bridge/src/context/impls.rs b/compiler/rustc_public_bridge/src/context/impls.rs index 87a8661edeb65..a018187e943e8 100644 --- a/compiler/rustc_public_bridge/src/context/impls.rs +++ b/compiler/rustc_public_bridge/src/context/impls.rs @@ -485,6 +485,16 @@ impl<'tcx, B: Bridge> CompilerCtxt<'tcx, B> { ty::Const::zero_sized(self.tcx, ty_internal) } + /// Create a caller location constant from a span. + /// + /// This produces a `&'static core::panic::Location<'static>` constant, + /// which is the implicit extra argument for `#[track_caller]` functions. + pub fn span_as_caller_location(&self, span: Span) -> MirConst<'tcx> { + let val = self.tcx.span_as_caller_location(span); + let ty = self.tcx.caller_location_ty(); + MirConst::from_value(val, ty) + } + /// Create a new constant that represents the given string value. pub fn new_const_str(&self, value: &str) -> MirConst<'tcx> { let ty = Ty::new_static_str(self.tcx); @@ -648,6 +658,15 @@ impl<'tcx, B: Bridge> CompilerCtxt<'tcx, B> { matches!(instance.def, ty::InstanceKind::Shim(ty::ShimKind::DropGlue(_, None))) } + /// Check if this instance requires a caller location argument. + /// + /// Functions with `#[track_caller]` have an implicit extra + /// `&'static core::panic::Location<'static>` argument appended to their ABI, + /// which is not visible in their MIR body signature. + pub fn instance_requires_caller_location(&self, instance: ty::Instance<'tcx>) -> bool { + instance.def.requires_caller_location(self.tcx) + } + /// Convert a non-generic crate item into an instance. /// This function will panic if the item is generic. pub fn mono_instance(&self, def_id: DefId) -> Instance<'tcx> { diff --git a/tests/ui-fulldeps/rustc_public/check_track_caller.rs b/tests/ui-fulldeps/rustc_public/check_track_caller.rs new file mode 100644 index 0000000000000..a9d396a479fa5 --- /dev/null +++ b/tests/ui-fulldeps/rustc_public/check_track_caller.rs @@ -0,0 +1,164 @@ +//@ run-pass +//! Test `#[track_caller]` support in rustc_public: +//! - `Instance::requires_caller_location` detects the implicit extra argument. +//! - `Body::caller_location` resolves the correct location constant. + +//@ ignore-stage1 +//@ ignore-cross-compile +//@ ignore-remote +//@ edition: 2021 + +#![feature(rustc_private)] + +extern crate rustc_driver; +extern crate rustc_interface; +extern crate rustc_middle; +#[macro_use] +extern crate rustc_public; + +use std::io::Write; +use std::ops::ControlFlow; + +use rustc_public::crate_def::CrateDef; +use rustc_public::mir::mono::Instance; +use rustc_public::mir::{Body, TerminatorKind}; +use rustc_public::ty::{ConstantKind, MirConst, RigidTy, TyKind}; + +const CRATE_NAME: &str = "input"; + +fn test_track_caller() -> ControlFlow<()> { + let items = rustc_public::all_local_items(); + + test_requires_caller_location(&items); + test_caller_location_from_regular_fn(&items); + test_caller_location_propagation(&items); + + ControlFlow::Continue(()) +} + +/// Check `requires_caller_location` and the extra ABI argument. +fn test_requires_caller_location(items: &[rustc_public::CrateItem]) { + let tracked = + items.iter().find(|item| item.name() == "input::tracked_fn").expect("missing tracked_fn"); + let not_tracked = items + .iter() + .find(|item| item.name() == "input::not_tracked_fn") + .expect("missing not_tracked_fn"); + + let tracked_instance = Instance::try_from(*tracked).unwrap(); + let not_tracked_instance = Instance::try_from(*not_tracked).unwrap(); + + assert!(tracked_instance.requires_caller_location()); + assert!(!not_tracked_instance.requires_caller_location()); + + // tracked_fn(u32) -> u32: MIR has 1 arg, ABI has 2 (extra &Location) + let tracked_abi = tracked_instance.fn_abi().unwrap(); + assert_eq!(tracked_abi.args.len(), 2, "tracked_fn ABI should have 2 args"); + + // not_tracked_fn(u32) -> u32: MIR has 1 arg, ABI has 1 + let not_tracked_abi = not_tracked_instance.fn_abi().unwrap(); + assert_eq!(not_tracked_abi.args.len(), 1, "not_tracked_fn ABI should have 1 arg"); +} + +/// Check that `Body::caller_location` resolves from the call site span for regular functions. +fn test_caller_location_from_regular_fn(items: &[rustc_public::CrateItem]) { + let caller = items.iter().find(|item| item.name() == "input::caller").expect("missing caller"); + let caller_instance = Instance::try_from(*caller).unwrap(); + assert!(!caller_instance.requires_caller_location()); + + let body = caller_instance.body().unwrap(); + let location = resolve_tracked_call_location(&body, None); + + // The result should be a &'static Location<'static> constant. + assert!( + matches!(location.ty().kind(), TyKind::RigidTy(RigidTy::Ref(..))), + "caller_location should produce a reference type" + ); + assert!( + matches!(location.kind(), ConstantKind::Allocated(..)), + "caller_location should produce an allocated constant" + ); +} + +/// Check that `Body::caller_location` propagates the inherited location for `#[track_caller]` fns. +fn test_caller_location_propagation(items: &[rustc_public::CrateItem]) { + let caller = items.iter().find(|item| item.name() == "input::caller").expect("missing caller"); + let caller_instance = Instance::try_from(*caller).unwrap(); + let caller_body = caller_instance.body().unwrap(); + + let wrapper = items + .iter() + .find(|item| item.name() == "input::tracked_wrapper") + .expect("missing tracked_wrapper"); + let wrapper_instance = Instance::try_from(*wrapper).unwrap(); + assert!(wrapper_instance.requires_caller_location()); + + let wrapper_body = wrapper_instance.body().unwrap(); + + // Use the location resolved from caller() as the inherited value. + let inherited = resolve_tracked_call_location(&caller_body, None); + + // tracked_wrapper is #[track_caller], so pass the inherited location. + // Without inlining, the inherited value should be returned as-is. + let result = resolve_tracked_call_location(&wrapper_body, Some(inherited.clone())); + assert_eq!(result, inherited, "inherited location should be propagated through"); +} + +/// Find the first call to a `#[track_caller]` function in the body and resolve its location. +fn resolve_tracked_call_location(body: &Body, inherited: Option) -> MirConst { + for bb in &body.blocks { + if let TerminatorKind::Call { func, .. } = &bb.terminator.kind { + let TyKind::RigidTy(RigidTy::FnDef(def, args)) = + func.ty(body.locals()).unwrap().kind() + else { + continue; + }; + let callee = Instance::resolve(def, &args).unwrap(); + if callee.requires_caller_location() { + return body.caller_location(&bb.terminator, inherited); + } + } + } + panic!("no call to a #[track_caller] function found in body"); +} + +fn main() { + let path = "track_caller_input.rs"; + generate_input(&path).unwrap(); + let args = &[ + "rustc".to_string(), + "-Cpanic=abort".to_string(), + "--crate-type=lib".to_string(), + "--crate-name".to_string(), + CRATE_NAME.to_string(), + path.to_string(), + ]; + run!(args, test_track_caller).unwrap(); +} + +fn generate_input(path: &str) -> std::io::Result<()> { + let mut file = std::fs::File::create(path)?; + write!( + file, + r#" + #[track_caller] + pub fn tracked_fn(x: u32) -> u32 {{ + x + 1 + }} + + pub fn not_tracked_fn(x: u32) -> u32 {{ + x + 1 + }} + + pub fn caller() -> u32 {{ + tracked_fn(42) + }} + + #[track_caller] + pub fn tracked_wrapper(x: u32) -> u32 {{ + tracked_fn(x) + }} + "# + )?; + Ok(()) +}