Implement #[track_caller]

Fixes rust-lang#848

Author: bjorn3

example/track-caller-attribute.rs

@@ -0,0 +1,42 @@
+// Based on https://github.com/anp/rust/blob/175631311716d7dfeceec40d2587cde7142ffa8c/src/test/ui/rfc-2091-track-caller/track-caller-attribute.rs
+
+// run-pass
+
+#![feature(track_caller)]
+
+use std::panic::Location;
+
+#[track_caller]
+fn tracked() -> &'static Location<'static> {
+    Location::caller()
+}
+
+fn nested_intrinsic() -> &'static Location<'static> {
+    Location::caller()
+}
+
+fn nested_tracked() -> &'static Location<'static> {
+    tracked()
+}
+
+fn main() {
+    let location = Location::caller();
+    assert_eq!(location.file(), file!());
+    assert_eq!(location.line(), 23);
+    assert_eq!(location.column(), 20);
+
+    let tracked = tracked();
+    assert_eq!(tracked.file(), file!());
+    assert_eq!(tracked.line(), 28);
+    assert_eq!(tracked.column(), 19);
+
+    let nested = nested_intrinsic();
+    assert_eq!(nested.file(), file!());
+    assert_eq!(nested.line(), 15);
+    assert_eq!(nested.column(), 5);
+
+    let contained = nested_tracked();
+    assert_eq!(contained.file(), file!());
+    assert_eq!(contained.line(), 19);
+    assert_eq!(contained.column(), 5);
+}

src/abi/comments.rs

@@ -14,27 +14,34 @@ pub fn add_args_header_comment(fx: &mut FunctionCx<impl Backend>) {
 pub fn add_arg_comment<'tcx>(
     fx: &mut FunctionCx<'_, 'tcx, impl Backend>,
     kind: &str,
-    local: mir::Local,
+    local: Option<mir::Local>,
     local_field: Option<usize>,
     params: EmptySinglePair<Value>,
     pass_mode: PassMode,
     ty: Ty<'tcx>,
 ) {
+    let local = if let Some(local) = local {
+        Cow::Owned(format!("{:?}", local))
+    } else {
+        Cow::Borrowed("???")
+    };
     let local_field = if let Some(local_field) = local_field {
         Cow::Owned(format!(".{}", local_field))
     } else {
         Cow::Borrowed("")
     };
+
     let params = match params {
         Empty => Cow::Borrowed("-"),
         Single(param) => Cow::Owned(format!("= {:?}", param)),
         Pair(param_a, param_b) => Cow::Owned(format!("= {:?}, {:?}", param_a, param_b)),
     };
+
     let pass_mode = format!("{:?}", pass_mode);
     fx.add_global_comment(format!(
         "{kind:5}{local:>3}{local_field:<5} {params:10} {pass_mode:36} {ty:?}",
         kind = kind,
-        local = format!("{:?}", local),
+        local = local,
         local_field = local_field,
         params = params,
         pass_mode = pass_mode,

src/abi/mod.rs

@@ -80,6 +80,7 @@ fn clif_sig_from_fn_sig<'tcx>(
     triple: &target_lexicon::Triple,
     sig: FnSig<'tcx>,
     is_vtable_fn: bool,
+    requires_caller_location: bool,
 ) -> Signature {
     let abi = match sig.abi {
         Abi::System => {
@@ -125,7 +126,7 @@ fn clif_sig_from_fn_sig<'tcx>(
         })
         .flatten();
 
-    let (params, returns) = match get_pass_mode(
+    let (mut params, returns): (Vec<_>, Vec<_>) = match get_pass_mode(
         tcx,
         tcx.layout_of(ParamEnv::reveal_all().and(output)).unwrap(),
     ) {
@@ -150,6 +151,10 @@ fn clif_sig_from_fn_sig<'tcx>(
         }
     };
 
+    if requires_caller_location {
+        params.push(AbiParam::new(pointer_ty(tcx)));
+    }
+
     Signature {
         params,
         returns,
@@ -169,7 +174,7 @@ pub fn get_function_name_and_sig<'tcx>(
     if fn_sig.c_variadic && !support_vararg {
         unimpl!("Variadic function definitions are not yet supported");
     }
-    let sig = clif_sig_from_fn_sig(tcx, triple, fn_sig, false);
+    let sig = clif_sig_from_fn_sig(tcx, triple, fn_sig, false, inst.def.requires_caller_location(tcx));
     (tcx.symbol_name(inst).name.as_str().to_string(), sig)
 }
 
@@ -316,18 +321,24 @@ pub fn codegen_fn_prelude(fx: &mut FunctionCx<'_, '_, impl Backend>, start_ebb:
 
                 let mut params = Vec::new();
                 for (i, arg_ty) in tupled_arg_tys.types().enumerate() {
-                    let param = cvalue_for_param(fx, start_ebb, local, Some(i), arg_ty);
+                    let param = cvalue_for_param(fx, start_ebb, Some(local), Some(i), arg_ty);
                     params.push(param);
                 }
 
                 (local, ArgKind::Spread(params), arg_ty)
             } else {
-                let param = cvalue_for_param(fx, start_ebb, local, None, arg_ty);
+                let param = cvalue_for_param(fx, start_ebb, Some(local), None, arg_ty);
                 (local, ArgKind::Normal(param), arg_ty)
             }
         })
         .collect::<Vec<(Local, ArgKind, Ty)>>();
 
+    assert!(fx.caller_location.is_none());
+    if fx.instance.def.requires_caller_location(fx.tcx) {
+        // Store caller location for `#[track_caller]`.
+        fx.caller_location = Some(cvalue_for_param(fx, start_ebb, None, None, fx.tcx.caller_location_ty()).unwrap());
+    }
+
     fx.bcx.switch_to_block(start_ebb);
     fx.bcx.ins().nop();
 
@@ -403,10 +414,10 @@ pub fn codegen_fn_prelude(fx: &mut FunctionCx<'_, '_, impl Backend>, start_ebb:
 
 pub fn codegen_terminator_call<'tcx>(
     fx: &mut FunctionCx<'_, 'tcx, impl Backend>,
+    span: Span,
     func: &Operand<'tcx>,
     args: &[Operand<'tcx>],
     destination: &Option<(Place<'tcx>, BasicBlock)>,
-    span: Span,
 ) {
     let fn_ty = fx.monomorphize(&func.ty(fx.mir, fx.tcx));
     let sig = fx
@@ -472,6 +483,7 @@ pub fn codegen_terminator_call<'tcx>(
 
     codegen_call_inner(
         fx,
+        span,
         Some(func),
         fn_ty,
         args,
@@ -488,6 +500,7 @@ pub fn codegen_terminator_call<'tcx>(
 
 fn codegen_call_inner<'tcx>(
     fx: &mut FunctionCx<'_, 'tcx, impl Backend>,
+    span: Span,
     func: Option<&Operand<'tcx>>,
     fn_ty: Ty<'tcx>,
     args: Vec<CValue<'tcx>>,
@@ -558,7 +571,7 @@ fn codegen_call_inner<'tcx>(
 
     let (call_inst, call_args) =
         self::returning::codegen_with_call_return_arg(fx, fn_sig, ret_place, |fx, return_ptr| {
-            let call_args: Vec<Value> = return_ptr
+            let mut call_args: Vec<Value> = return_ptr
                 .into_iter()
                 .chain(first_arg.into_iter())
                 .chain(
@@ -569,9 +582,20 @@ fn codegen_call_inner<'tcx>(
                 )
                 .collect::<Vec<_>>();
 
+            if instance.map(|inst| inst.def.requires_caller_location(fx.tcx)).unwrap_or(false) {
+                // Pass the caller location for `#[track_caller]`.
+                let caller_location = fx.get_caller_location(span);
+                call_args.extend(adjust_arg_for_abi(fx, caller_location).into_iter());
+            }
+
             let call_inst = if let Some(func_ref) = func_ref {
-                let sig =
-                    clif_sig_from_fn_sig(fx.tcx, fx.triple(), fn_sig, is_virtual_call);
+                let sig = clif_sig_from_fn_sig(
+                    fx.tcx,
+                    fx.triple(),
+                    fn_sig,
+                    is_virtual_call,
+                    false, // calls through function pointers never pass the caller location
+                );
                 let sig = fx.bcx.import_signature(sig);
                 fx.bcx.ins().call_indirect(sig, func_ref, &call_args)
             } else {
@@ -604,7 +628,11 @@ fn codegen_call_inner<'tcx>(
     }
 }
 
-pub fn codegen_drop<'tcx>(fx: &mut FunctionCx<'_, 'tcx, impl Backend>, drop_place: CPlace<'tcx>) {
+pub fn codegen_drop<'tcx>(
+    fx: &mut FunctionCx<'_, 'tcx, impl Backend>,
+    span: Span,
+    drop_place: CPlace<'tcx>,
+) {
     let ty = drop_place.layout().ty;
     let drop_fn = Instance::resolve_drop_in_place(fx.tcx, ty);
 
@@ -625,7 +653,13 @@ pub fn codegen_drop<'tcx>(fx: &mut FunctionCx<'_, 'tcx, impl Backend>, drop_plac
 
                 assert_eq!(fn_sig.output(), fx.tcx.mk_unit());
 
-                let sig = clif_sig_from_fn_sig(fx.tcx, fx.triple(), fn_sig, true);
+                let sig = clif_sig_from_fn_sig(
+                    fx.tcx,
+                    fx.triple(),
+                    fn_sig,
+                    true,
+                    false, // `drop_in_place` is never `#[track_caller]`
+                );
                 let sig = fx.bcx.import_signature(sig);
                 fx.bcx.ins().call_indirect(sig, drop_fn, &[ptr]);
             }
@@ -642,7 +676,7 @@ pub fn codegen_drop<'tcx>(fx: &mut FunctionCx<'_, 'tcx, impl Backend>, drop_plac
                 );
                 drop_place.write_place_ref(fx, arg_place);
                 let arg_value = arg_place.to_cvalue(fx);
-                codegen_call_inner(fx, None, drop_fn_ty, vec![arg_value], None);
+                codegen_call_inner(fx, span, None, drop_fn_ty, vec![arg_value], None);
             }
         }
     }

src/abi/pass_mode.rs

@@ -126,7 +126,7 @@ pub fn adjust_arg_for_abi<'tcx>(
 pub fn cvalue_for_param<'tcx>(
     fx: &mut FunctionCx<'_, 'tcx, impl Backend>,
     start_ebb: Ebb,
-    local: mir::Local,
+    local: Option<mir::Local>,
     local_field: Option<usize>,
     arg_ty: Ty<'tcx>,
 ) -> Option<CValue<'tcx>> {

src/abi/returning.rs

@@ -46,7 +46,7 @@ pub fn codegen_return_param(
     crate::abi::comments::add_arg_comment(
         fx,
         "ret",
-        RETURN_PLACE,
+        Some(RETURN_PLACE),
         None,
         ret_param,
         ret_pass_mode,

src/base.rs

@@ -48,6 +48,7 @@ pub fn trans_fn<'clif, 'tcx, B: Backend + 'static>(
         bcx,
         ebb_map,
         local_map: HashMap::new(),
+        caller_location: None, // set by `codegen_fn_prelude`
 
         clif_comments,
         constants_cx: &mut cx.constants_cx,
@@ -236,10 +237,10 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, impl Backend>) {
             } => {
                 fx.tcx.sess.time("codegen call", || crate::abi::codegen_terminator_call(
                     fx,
+                    bb_data.terminator().source_info.span,
                     func,
                     args,
                     destination,
-                    bb_data.terminator().source_info.span,
                 ));
             }
             TerminatorKind::Resume | TerminatorKind::Abort => {
@@ -261,7 +262,7 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, impl Backend>) {
                 unwind: _,
             } => {
                 let drop_place = trans_place(fx, location);
-                crate::abi::codegen_drop(fx, drop_place);
+                crate::abi::codegen_drop(fx, bb_data.terminator().source_info.span, drop_place);
 
                 let target_ebb = fx.get_ebb(*target);
                 fx.bcx.ins().jump(target_ebb, &[]);
@@ -370,7 +371,7 @@ fn trans_stmt<'tcx>(
                     match from_ty.kind {
                         ty::FnDef(def_id, substs) => {
                             let func_ref = fx.get_function_ref(
-                                Instance::resolve(fx.tcx, ParamEnv::reveal_all(), def_id, substs)
+                                Instance::resolve_for_fn_ptr(fx.tcx, ParamEnv::reveal_all(), def_id, substs)
                                     .unwrap(),
                             );
                             let func_addr = fx.bcx.ins().func_addr(fx.pointer_type, func_ref);

src/common.rs

@@ -267,6 +267,9 @@ pub struct FunctionCx<'clif, 'tcx, B: Backend + 'static> {
     pub ebb_map: IndexVec<BasicBlock, Ebb>,
     pub local_map: HashMap<Local, CPlace<'tcx>>,
 
+    /// When `#[track_caller]` is used, the implicit caller location is stored in this variable.
+    pub caller_location: Option<CValue<'tcx>>,
+
     pub clif_comments: crate::pretty_clif::CommentWriter,
     pub constants_cx: &'clif mut crate::constant::ConstantCx,
     pub vtables: &'clif mut HashMap<(Ty<'tcx>, Option<ty::PolyExistentialTraitRef<'tcx>>), DataId>,
@@ -355,6 +358,11 @@ impl<'tcx, B: Backend + 'static> FunctionCx<'_, 'tcx, B> {
     }
 
     pub 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;
+        }
+
         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((

test.sh

@@ -64,6 +64,10 @@ echo "[AOT] subslice-patterns-const-eval"
 $RUSTC example/subslice-patterns-const-eval.rs --crate-type bin -Cpanic=abort
 ./target/out/subslice-patterns-const-eval
 
+echo "[AOT] track-caller-attribute"
+$RUSTC example/track-caller-attribute.rs --crate-type bin -Cpanic=abort
+./target/out/track-caller-attribute
+
 echo "[BUILD] mod_bench"
 $RUSTC example/mod_bench.rs --crate-type bin