Fix #[track_caller] with MIR inlining

Author: bjorn3

scripts/test_rustc_tests.sh

@@ -105,8 +105,6 @@ rm src/test/ui/type-alias-impl-trait/assoc-projection-ice.rs # produces ICE
 
 rm src/test/ui/simd/intrinsic/generic-reduction-pass.rs # simd_reduce_add_unordered doesn't accept an accumulator for integer vectors
 
-rm src/test/ui/rfc-2091-track-caller/intrinsic-wrapper.rs # wrong result from `Location::caller()`
-
 # bugs in the test suite
 # ======================
 rm src/test/ui/backtrace.rs # TODO warning

src/abi/mod.rs

@@ -309,7 +309,7 @@ fn codegen_call_argument_operand<'tcx>(
 
 pub(crate) fn codegen_terminator_call<'tcx>(
     fx: &mut FunctionCx<'_, '_, 'tcx>,
-    span: Span,
+    source_info: mir::SourceInfo,
     func: &Operand<'tcx>,
     args: &[Operand<'tcx>],
     mir_dest: Option<(Place<'tcx>, BasicBlock)>,
@@ -340,7 +340,13 @@ pub(crate) fn codegen_terminator_call<'tcx>(
 
         match instance.def {
             InstanceDef::Intrinsic(_) => {
-                crate::intrinsics::codegen_intrinsic_call(fx, instance, args, destination, span);
+                crate::intrinsics::codegen_intrinsic_call(
+                    fx,
+                    instance,
+                    args,
+                    destination,
+                    source_info,
+                );
                 return;
             }
             InstanceDef::DropGlue(_, None) => {
@@ -402,7 +408,7 @@ pub(crate) fn codegen_terminator_call<'tcx>(
 
     // Pass the caller location for `#[track_caller]`.
     if instance.map(|inst| inst.def.requires_caller_location(fx.tcx)).unwrap_or(false) {
-        let caller_location = fx.get_caller_location(span);
+        let caller_location = fx.get_caller_location(source_info);
         args.push(CallArgument { value: caller_location, is_owned: false });
     }
 
@@ -479,9 +485,10 @@ pub(crate) fn codegen_terminator_call<'tcx>(
         // FIXME find a cleaner way to support varargs
         if fn_sig.c_variadic {
             if !matches!(fn_sig.abi, Abi::C { .. }) {
-                fx.tcx
-                    .sess
-                    .span_fatal(span, &format!("Variadic call for non-C abi {:?}", fn_sig.abi));
+                fx.tcx.sess.span_fatal(
+                    source_info.span,
+                    &format!("Variadic call for non-C abi {:?}", fn_sig.abi),
+                );
             }
             let sig_ref = fx.bcx.func.dfg.call_signature(call_inst).unwrap();
             let abi_params = call_args
@@ -490,9 +497,10 @@ pub(crate) fn codegen_terminator_call<'tcx>(
                     let ty = fx.bcx.func.dfg.value_type(arg);
                     if !ty.is_int() {
                         // FIXME set %al to upperbound on float args once floats are supported
-                        fx.tcx
-                            .sess
-                            .span_fatal(span, &format!("Non int ty {:?} for variadic call", ty));
+                        fx.tcx.sess.span_fatal(
+                            source_info.span,
+                            &format!("Non int ty {:?} for variadic call", ty),
+                        );
                     }
                     AbiParam::new(ty)
                 })
@@ -513,7 +521,7 @@ pub(crate) fn codegen_terminator_call<'tcx>(
 
 pub(crate) fn codegen_drop<'tcx>(
     fx: &mut FunctionCx<'_, '_, 'tcx>,
-    span: Span,
+    source_info: mir::SourceInfo,
     drop_place: CPlace<'tcx>,
 ) {
     let ty = drop_place.layout().ty;
@@ -560,7 +568,7 @@ pub(crate) fn codegen_drop<'tcx>(
 
                 if drop_instance.def.requires_caller_location(fx.tcx) {
                     // Pass the caller location for `#[track_caller]`.
-                    let caller_location = fx.get_caller_location(span);
+                    let caller_location = fx.get_caller_location(source_info);
                     call_args.extend(
                         adjust_arg_for_abi(fx, caller_location, &fn_abi.args[1], false).into_iter(),
                     );

src/base.rs

@@ -325,7 +325,7 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, '_>) {
                     AssertKind::BoundsCheck { ref len, ref index } => {
                         let len = codegen_operand(fx, len).load_scalar(fx);
                         let index = codegen_operand(fx, index).load_scalar(fx);
-                        let location = fx.get_caller_location(source_info.span).load_scalar(fx);
+                        let location = fx.get_caller_location(source_info).load_scalar(fx);
 
                         codegen_panic_inner(
                             fx,
@@ -336,7 +336,7 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, '_>) {
                     }
                     _ => {
                         let msg_str = msg.description();
-                        codegen_panic(fx, msg_str, source_info.span);
+                        codegen_panic(fx, msg_str, source_info);
                     }
                 }
             }
@@ -393,12 +393,12 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, '_>) {
                 func,
                 args,
                 destination,
-                fn_span,
+                fn_span: _,
                 cleanup: _,
                 from_hir_call: _,
             } => {
                 fx.tcx.sess.time("codegen call", || {
-                    crate::abi::codegen_terminator_call(fx, *fn_span, func, args, *destination)
+                    crate::abi::codegen_terminator_call(fx, source_info, func, args, *destination)
                 });
             }
             TerminatorKind::InlineAsm {
@@ -450,7 +450,7 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, '_>) {
             }
             TerminatorKind::Drop { place, target, unwind: _ } => {
                 let drop_place = codegen_place(fx, *place);
-                crate::abi::codegen_drop(fx, source_info.span, drop_place);
+                crate::abi::codegen_drop(fx, source_info, drop_place);
 
                 let target_block = fx.get_block(*target);
                 fx.bcx.ins().jump(target_block, &[]);
@@ -898,14 +898,18 @@ pub(crate) fn codegen_operand<'tcx>(
     }
 }
 
-pub(crate) fn codegen_panic<'tcx>(fx: &mut FunctionCx<'_, '_, 'tcx>, msg_str: &str, span: Span) {
-    let location = fx.get_caller_location(span).load_scalar(fx);
+pub(crate) fn codegen_panic<'tcx>(
+    fx: &mut FunctionCx<'_, '_, 'tcx>,
+    msg_str: &str,
+    source_info: mir::SourceInfo,
+) {
+    let location = fx.get_caller_location(source_info).load_scalar(fx);
 
     let msg_ptr = fx.anonymous_str(msg_str);
     let msg_len = fx.bcx.ins().iconst(fx.pointer_type, i64::try_from(msg_str.len()).unwrap());
     let args = [msg_ptr, msg_len, location];
 
-    codegen_panic_inner(fx, rustc_hir::LangItem::Panic, &args, span);
+    codegen_panic_inner(fx, rustc_hir::LangItem::Panic, &args, source_info.span);
 }
 
 pub(crate) fn codegen_panic_inner<'tcx>(

src/common.rs

@@ -340,22 +340,46 @@ impl<'tcx> FunctionCx<'_, '_, 'tcx> {
         self.bcx.set_srcloc(SourceLoc::new(index as u32));
     }
 
-    pub(crate) fn get_caller_location(&mut self, span: Span) -> CValue<'tcx> {
-        if let Some(loc) = self.caller_location {
-            // `#[track_caller]` is used; return caller location instead of current location.
-            return loc;
+    // Note: must be kept in sync with get_caller_location from cg_ssa
+    pub(crate) fn get_caller_location(&mut self, mut source_info: mir::SourceInfo) -> CValue<'tcx> {
+        let span_to_caller_location = |fx: &mut FunctionCx<'_, '_, 'tcx>, span: Span| {
+            let topmost = span.ctxt().outer_expn().expansion_cause().unwrap_or(span);
+            let caller = fx.tcx.sess.source_map().lookup_char_pos(topmost.lo());
+            let const_loc = fx.tcx.const_caller_location((
+                rustc_span::symbol::Symbol::intern(
+                    &caller.file.name.prefer_remapped().to_string_lossy(),
+                ),
+                caller.line as u32,
+                caller.col_display as u32 + 1,
+            ));
+            crate::constant::codegen_const_value(fx, const_loc, fx.tcx.caller_location_ty())
+        };
+
+        // Walk up the `SourceScope`s, in case some of them are from MIR inlining.
+        // If so, the starting `source_info.span` is in the innermost inlined
+        // function, and will be replaced with outer callsite spans as long
+        // as the inlined functions were `#[track_caller]`.
+        loop {
+            let scope_data = &self.mir.source_scopes[source_info.scope];
+
+            if let Some((callee, callsite_span)) = scope_data.inlined {
+                // Stop inside the most nested non-`#[track_caller]` function,
+                // before ever reaching its caller (which is irrelevant).
+                if !callee.def.requires_caller_location(self.tcx) {
+                    return span_to_caller_location(self, source_info.span);
+                }
+                source_info.span = callsite_span;
+            }
+
+            // Skip past all of the parents with `inlined: None`.
+            match scope_data.inlined_parent_scope {
+                Some(parent) => source_info.scope = parent,
+                None => break,
+            }
         }
 
-        let topmost = span.ctxt().outer_expn().expansion_cause().unwrap_or(span);
-        let caller = self.tcx.sess.source_map().lookup_char_pos(topmost.lo());
-        let const_loc = self.tcx.const_caller_location((
-            rustc_span::symbol::Symbol::intern(
-                &caller.file.name.prefer_remapped().to_string_lossy(),
-            ),
-            caller.line as u32,
-            caller.col_display as u32 + 1,
-        ));
-        crate::constant::codegen_const_value(self, const_loc, self.tcx.caller_location_ty())
+        // No inlined `SourceScope`s, or all of them were `#[track_caller]`.
+        self.caller_location.unwrap_or_else(|| span_to_caller_location(self, source_info.span))
     }
 
     pub(crate) fn anonymous_str(&mut self, msg: &str) -> Value {