rustc_mir: support MIR-inlining #[track_caller] functions.

Author: eddyb

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -405,7 +405,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         self.set_debug_loc(&mut bx, terminator.source_info);
 
         // Get the location information.
-        let location = self.get_caller_location(&mut bx, span).immediate();
+        let location = self.get_caller_location(&mut bx, terminator.source_info).immediate();
 
         // Put together the arguments to the panic entry point.
         let (lang_item, args) = match msg {
@@ -442,7 +442,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         bx: &mut Bx,
         intrinsic: Option<Symbol>,
         instance: Option<Instance<'tcx>>,
-        span: Span,
+        source_info: mir::SourceInfo,
         destination: &Option<(mir::Place<'tcx>, mir::BasicBlock)>,
         cleanup: Option<mir::BasicBlock>,
     ) -> bool {
@@ -484,11 +484,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     }
                 });
                 let msg = bx.const_str(Symbol::intern(&msg_str));
-                let location = self.get_caller_location(bx, span).immediate();
+                let location = self.get_caller_location(bx, source_info).immediate();
 
                 // Obtain the panic entry point.
                 // FIXME: dedup this with `codegen_assert_terminator` above.
-                let def_id = common::langcall(bx.tcx(), Some(span), "", LangItem::Panic);
+                let def_id =
+                    common::langcall(bx.tcx(), Some(source_info.span), "", LangItem::Panic);
                 let instance = ty::Instance::mono(bx.tcx(), def_id);
                 let fn_abi = FnAbi::of_instance(bx, instance, &[]);
                 let llfn = bx.get_fn_addr(instance);
@@ -529,7 +530,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         cleanup: Option<mir::BasicBlock>,
         fn_span: Span,
     ) {
-        let span = terminator.source_info.span;
+        let source_info = terminator.source_info;
+        let span = source_info.span;
+
         // Create the callee. This is a fn ptr or zero-sized and hence a kind of scalar.
         let callee = self.codegen_operand(&mut bx, func);
 
@@ -606,7 +609,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             &mut bx,
             intrinsic,
             instance,
-            span,
+            source_info,
             destination,
             cleanup,
         ) {
@@ -627,7 +630,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
         if intrinsic == Some(sym::caller_location) {
             if let Some((_, target)) = destination.as_ref() {
-                let location = self.get_caller_location(&mut bx, fn_span);
+                let location = self
+                    .get_caller_location(&mut bx, mir::SourceInfo { span: fn_span, ..source_info });
 
                 if let ReturnDest::IndirectOperand(tmp, _) = ret_dest {
                     location.val.store(&mut bx, tmp);
@@ -686,7 +690,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 &fn_abi,
                 &args,
                 dest,
-                terminator.source_info.span,
+                span,
             );
 
             if let ReturnDest::IndirectOperand(dst, _) = ret_dest {
@@ -793,7 +797,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 args.len() + 1,
                 "#[track_caller] fn's must have 1 more argument in their ABI than in their MIR",
             );
-            let location = self.get_caller_location(&mut bx, fn_span);
+            let location =
+                self.get_caller_location(&mut bx, mir::SourceInfo { span: fn_span, ..source_info });
             debug!(
                 "codegen_call_terminator({:?}): location={:?} (fn_span {:?})",
                 terminator, location, fn_span
@@ -1179,17 +1184,49 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
     }
 
-    fn get_caller_location(&mut self, bx: &mut Bx, span: Span) -> OperandRef<'tcx, Bx::Value> {
-        self.caller_location.unwrap_or_else(|| {
+    fn get_caller_location(
+        &mut self,
+        bx: &mut Bx,
+        mut source_info: mir::SourceInfo,
+    ) -> OperandRef<'tcx, Bx::Value> {
+        let tcx = bx.tcx();
+
+        let mut span_to_caller_location = |span: Span| {
             let topmost = span.ctxt().outer_expn().expansion_cause().unwrap_or(span);
-            let caller = bx.tcx().sess.source_map().lookup_char_pos(topmost.lo());
-            let const_loc = bx.tcx().const_caller_location((
+            let caller = tcx.sess.source_map().lookup_char_pos(topmost.lo());
+            let const_loc = tcx.const_caller_location((
                 Symbol::intern(&caller.file.name.to_string()),
                 caller.line as u32,
                 caller.col_display as u32 + 1,
             ));
             OperandRef::from_const(bx, const_loc, bx.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(tcx) {
+                    return span_to_caller_location(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,
+            }
+        }
+
+        // No inlined `SourceScope`s, or all of them were `#[track_caller]`.
+        self.caller_location.unwrap_or_else(|| span_to_caller_location(source_info.span))
     }
 
     fn get_personality_slot(&mut self, bx: &mut Bx) -> PlaceRef<'tcx, Bx::Value> {

compiler/rustc_middle/src/mir/mod.rs

@@ -1878,6 +1878,11 @@ pub struct SourceScopeData<'tcx> {
     /// `ty::Instance` is the callee, and the `Span` is the call site.
     pub inlined: Option<(ty::Instance<'tcx>, Span)>,
 
+    /// Nearest (transitive) parent scope (if any) which is inlined.
+    /// This is an optimization over walking up `parent_scope`
+    /// until a scope with `inlined: Some(...)` is found.
+    pub inlined_parent_scope: Option<SourceScope>,
+
     /// Crate-local information for this source scope, that can't (and
     /// needn't) be tracked across crates.
     pub local_data: ClearCrossCrate<SourceScopeLocalData>,

compiler/rustc_middle/src/mir/visit.rs

@@ -325,6 +325,7 @@ macro_rules! make_mir_visitor {
                     span,
                     parent_scope,
                     inlined,
+                    inlined_parent_scope,
                     local_data: _,
                 } = scope_data;
 
@@ -357,6 +358,9 @@ macro_rules! make_mir_visitor {
                     }
                     self.visit_substs(callee_substs, location);
                 }
+                if let Some(inlined_parent_scope) = inlined_parent_scope {
+                    self.visit_source_scope(inlined_parent_scope);
+                }
             }
 
             fn super_statement(&mut self,

compiler/rustc_mir/src/interpret/intrinsics/caller_location.rs

@@ -15,38 +15,61 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     /// Walks up the callstack from the intrinsic's callsite, searching for the first callsite in a
     /// frame which is not `#[track_caller]`.
     crate fn find_closest_untracked_caller_location(&self) -> Span {
-        let frame = self
-            .stack()
-            .iter()
-            .rev()
-            // Find first non-`#[track_caller]` frame.
-            .find(|frame| {
+        for frame in self.stack().iter().rev() {
+            debug!("find_closest_untracked_caller_location: checking frame {:?}", frame.instance);
+
+            // Assert that the frame we look at is actually executing code currently
+            // (`loc` is `Err` when we are unwinding and the frame does not require cleanup).
+            let loc = frame.loc.unwrap();
+
+            // This could be a non-`Call` terminator (such as `Drop`), or not a terminator at all
+            // (such as `box`). Use the normal span by default.
+            let mut source_info = *frame.body.source_info(loc);
+
+            // If this is a `Call` terminator, use the `fn_span` instead.
+            let block = &frame.body.basic_blocks()[loc.block];
+            if loc.statement_index == block.statements.len() {
                 debug!(
-                    "find_closest_untracked_caller_location: checking frame {:?}",
-                    frame.instance
+                    "find_closest_untracked_caller_location: got terminator {:?} ({:?})",
+                    block.terminator(),
+                    block.terminator().kind
                 );
-                !frame.instance.def.requires_caller_location(*self.tcx)
-            })
-            // Assert that there is always such a frame.
-            .unwrap();
-        // Assert that the frame we look at is actually executing code currently
-        // (`loc` is `Err` when we are unwinding and the frame does not require cleanup).
-        let loc = frame.loc.unwrap();
-        // If this is a `Call` terminator, use the `fn_span` instead.
-        let block = &frame.body.basic_blocks()[loc.block];
-        if loc.statement_index == block.statements.len() {
-            debug!(
-                "find_closest_untracked_caller_location:: got terminator {:?} ({:?})",
-                block.terminator(),
-                block.terminator().kind
-            );
-            if let TerminatorKind::Call { fn_span, .. } = block.terminator().kind {
-                return fn_span;
+                if let TerminatorKind::Call { fn_span, .. } = block.terminator().kind {
+                    source_info.span = fn_span;
+                }
+            }
+
+            // 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 = &frame.body.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 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,
+                }
+            }
+
+            // Stop inside the most nested non-`#[track_caller]` function,
+            // before ever reaching its caller (which is irrelevant).
+            if !frame.instance.def.requires_caller_location(*self.tcx) {
+                return source_info.span;
             }
         }
-        // This is a different terminator (such as `Drop`) or not a terminator at all
-        // (such as `box`). Use the normal span.
-        frame.body.source_info(loc).span
+
+        bug!("no non-`#[track_caller]` frame found")
     }
 
     /// Allocate a `const core::panic::Location` with the provided filename and line/column numbers.

compiler/rustc_mir/src/shim.rs

@@ -216,6 +216,7 @@ fn new_body<'tcx>(
                 span,
                 parent_scope: None,
                 inlined: None,
+                inlined_parent_scope: None,
                 local_data: ClearCrossCrate::Clear,
             },
             1,

compiler/rustc_mir/src/transform/inline.rs

@@ -246,11 +246,6 @@ impl Inliner<'tcx> {
 
         let codegen_fn_attrs = tcx.codegen_fn_attrs(callsite.callee.def_id());
 
-        if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::TRACK_CALLER) {
-            debug!("`#[track_caller]` present - not inlining");
-            return false;
-        }
-
         let self_features = &self.codegen_fn_attrs.target_features;
         let callee_features = &codegen_fn_attrs.target_features;
         if callee_features.iter().any(|feature| !self_features.contains(feature)) {
@@ -441,11 +436,24 @@ impl Inliner<'tcx> {
 
                 for mut scope in callee_body.source_scopes.iter().cloned() {
                     if scope.parent_scope.is_none() {
+                        let callsite_scope = &caller_body.source_scopes[callsite.source_info.scope];
+
+                        // Attach the outermost callee scope as a child of the callsite
+                        // scope, via the `parent_scope` and `inlined_parent_scope` chains.
                         scope.parent_scope = Some(callsite.source_info.scope);
+                        assert_eq!(scope.inlined_parent_scope, None);
+                        scope.inlined_parent_scope = if callsite_scope.inlined.is_some() {
+                            Some(callsite.source_info.scope)
+                        } else {
+                            callsite_scope.inlined_parent_scope
+                        };
 
                         // Mark the outermost callee scope as an inlined one.
                         assert_eq!(scope.inlined, None);
                         scope.inlined = Some((callsite.callee, callsite.source_info.span));
+                    } else if scope.inlined_parent_scope.is_none() {
+                        // Make it easy to find the scope with `inlined` set above.
+                        scope.inlined_parent_scope = Some(scope_map[OUTERMOST_SOURCE_SCOPE]);
                     }
 
                     let idx = caller_body.source_scopes.push(scope);

compiler/rustc_mir/src/util/pretty.rs

@@ -551,18 +551,31 @@ fn write_scope_tree(
         let child_data = &body.source_scopes[child];
         assert_eq!(child_data.parent_scope, Some(parent));
 
-        if let Some((callee, callsite_span)) = child_data.inlined {
-            let indented_header =
-                format!("{0:1$}scope {2} (inlined {3}) {{", "", indent, child.index(), callee);
+        let (special, span) = if let Some((callee, callsite_span)) = child_data.inlined {
+            (
+                format!(
+                    " (inlined {}{})",
+                    if callee.def.requires_caller_location(tcx) { "#[track_caller] " } else { "" },
+                    callee
+                ),
+                Some(callsite_span),
+            )
+        } else {
+            (String::new(), None)
+        };
+
+        let indented_header = format!("{0:1$}scope {2}{3} {{", "", indent, child.index(), special);
+
+        if let Some(span) = span {
             writeln!(
                 w,
                 "{0:1$} // at {2}",
                 indented_header,
                 ALIGN,
-                tcx.sess.source_map().span_to_string(callsite_span),
+                tcx.sess.source_map().span_to_string(span),
             )?;
         } else {
-            writeln!(w, "{0:1$}scope {2} {{", "", indent, child.index())?;
+            writeln!(w, "{}", indented_header)?;
         }
 
         write_scope_tree(tcx, body, scope_tree, w, child, depth + 1)?;

compiler/rustc_mir_build/src/build/scope.rs

@@ -706,6 +706,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             span,
             parent_scope: Some(parent),
             inlined: None,
+            inlined_parent_scope: None,
             local_data: ClearCrossCrate::Set(scope_local_data),
         })
     }