Add support to caller_location to rustc_public

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| {

compiler/rustc_public/src/mir/body.rs

@@ -20,30 +20,40 @@ 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<VarDebugInfo>,
 
     /// 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<Local>,
+    pub(crate) spread_arg: Option<Local>,
 
     /// 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<SourceScopeInfo>,
 }
 
 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<BasicBlock>,
         locals: LocalDecls,
@@ -52,13 +62,15 @@ impl Body {
         spread_arg: Option<Local>,
         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<Local> {
         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>,
+    ) -> 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<LocalDecl>;
@@ -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<SourceScope>,
+}
+
 #[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)
+}

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

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);

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)]

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,19 +35,34 @@ 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),
                     span: decl.source_info.span.stable(tables, cx),
                     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(),
+        }
     }
 }
 

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> {