Support #[track_caller] on closures and generators

This PR allows applying a `#[track_caller]` attribute to a
closure/generator expression. The attribute as interpreted as applying
to the compiler-generated implementation of the corresponding trait
method (`FnOnce::call_once`, `FnMut::call_mut`, `Fn::call`, or
`Generator::resume`).

This feature does not have its own feature gate - however, it requires
`#![feature(stmt_expr_attributes)]` in order to actually apply
an attribute to a closure or generator.

This is implemented in the same way as for functions - an extra
location argument is appended to the end of the ABI. For closures,
this argument is *not* part of the 'tupled' argument storing the
parameters - the final closure argument for `#[track_caller]` closures
is no longer a tuple.

For direct (monomorphized) calls, the necessary support was already
implemented - we just needeed to adjust some assertions around checking
the ABI and argument count to take closures into account.

For calls through a trait object, more work was needed.
When creating a `ReifyShim`, we need to create a shim
for the trait method (e.g. `FnOnce::call_mut`) - unlike normal
functions, closures are never invoked directly, and always go through a
trait method.

Additional handling was needed for `InstanceDef::ClosureOnceShim`. In
order to pass location information throgh a direct (monomorphized) call
to `FnOnce::call_once` on an `FnMut` closure, we need to make
`ClosureOnceShim` aware of `#[tracked_caller]`. A new field
`track_caller` is added to `ClosureOnceShim` - this is used by
`InstanceDef::requires_caller` location, allowing codegen to
pass through the extra location argument.

Since `ClosureOnceShim.track_caller` is only used by codegen,
we end up generating two identical MIR shims - one for
`track_caller == true`, and one for `track_caller == false`. However,
these two shims are used by the entire crate (i.e. it's two shims total,
not two shims per unique closure), so this shouldn't a big deal.

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -777,22 +777,30 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
             self.codegen_argument(&mut bx, op, &mut llargs, &fn_abi.args[i]);
         }
-        if let Some(tup) = untuple {
+        let num_untupled = untuple.map(|tup| {
             self.codegen_arguments_untupled(
                 &mut bx,
                 tup,
                 &mut llargs,
                 &fn_abi.args[first_args.len()..],
             )
-        }
+        });
 
         let needs_location =
             instance.map_or(false, |i| i.def.requires_caller_location(self.cx.tcx()));
         if needs_location {
+            let mir_args = if let Some(num_untupled) = num_untupled {
+                first_args.len() + num_untupled
+            } else {
+                args.len()
+            };
             assert_eq!(
                 fn_abi.args.len(),
-                args.len() + 1,
-                "#[track_caller] fn's must have 1 more argument in their ABI than in their MIR",
+                mir_args + 1,
+                "#[track_caller] fn's must have 1 more argument in their ABI than in their MIR: {:?} {:?} {:?}",
+                instance,
+                fn_span,
+                fn_abi,
             );
             let location =
                 self.get_caller_location(&mut bx, mir::SourceInfo { span: fn_span, ..source_info });
@@ -1122,7 +1130,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         operand: &mir::Operand<'tcx>,
         llargs: &mut Vec<Bx::Value>,
         args: &[ArgAbi<'tcx, Ty<'tcx>>],
-    ) {
+    ) -> usize {
         let tuple = self.codegen_operand(bx, operand);
 
         // Handle both by-ref and immediate tuples.
@@ -1142,6 +1150,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 self.codegen_argument(bx, op, llargs, &args[i]);
             }
         }
+        tuple.layout.fields.count()
     }
 
     fn get_caller_location(

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -258,6 +258,8 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     let mut idx = 0;
     let mut llarg_idx = fx.fn_abi.ret.is_indirect() as usize;
 
+    let mut num_untupled = None;
+
     let args = mir
         .args_iter()
         .enumerate()
@@ -286,6 +288,11 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
                     let pr_field = place.project_field(bx, i);
                     bx.store_fn_arg(arg, &mut llarg_idx, pr_field);
                 }
+                assert_eq!(
+                    None,
+                    num_untupled.replace(tupled_arg_tys.len()),
+                    "Replaced existing num_tupled"
+                );
 
                 return LocalRef::Place(place);
             }
@@ -362,10 +369,17 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         .collect::<Vec<_>>();
 
     if fx.instance.def.requires_caller_location(bx.tcx()) {
+        let mir_args = if let Some(num_untupled) = num_untupled {
+            // Subtract off the tupled argument that gets 'expanded'
+            args.len() - 1 + num_untupled
+        } else {
+            args.len()
+        };
         assert_eq!(
             fx.fn_abi.args.len(),
-            args.len() + 1,
-            "#[track_caller] fn's must have 1 more argument in their ABI than in their MIR",
+            mir_args + 1,
+            "#[track_caller] instance {:?} must have 1 more argument in their ABI than in their MIR",
+            fx.instance
         );
 
         let arg = fx.fn_abi.args.last().unwrap();

compiler/rustc_feature/src/active.rs

@@ -678,6 +678,8 @@ declare_features! (
     /// Allows the `#[must_not_suspend]` attribute.
     (active, must_not_suspend, "1.57.0", Some(83310), None),
 
+    /// Allows `#[track_caller]` on closures and generators.
+    (active, closure_track_caller, "1.57.0", Some(87417), None),
 
     // -------------------------------------------------------------------------
     // feature-group-end: actual feature gates

compiler/rustc_middle/src/mir/visit.rs

@@ -348,7 +348,7 @@ macro_rules! make_mir_visitor {
                         ty::InstanceDef::VtableShim(_def_id) |
                         ty::InstanceDef::ReifyShim(_def_id) |
                         ty::InstanceDef::Virtual(_def_id, _) |
-                        ty::InstanceDef::ClosureOnceShim { call_once: _def_id } |
+                        ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |
                         ty::InstanceDef::DropGlue(_def_id, None) => {}
 
                         ty::InstanceDef::FnPtrShim(_def_id, ty) |

compiler/rustc_middle/src/ty/instance.rs

@@ -77,7 +77,7 @@ pub enum InstanceDef<'tcx> {
     /// `<[FnMut closure] as FnOnce>::call_once`.
     ///
     /// The `DefId` is the ID of the `call_once` method in `FnOnce`.
-    ClosureOnceShim { call_once: DefId },
+    ClosureOnceShim { call_once: DefId, track_caller: bool },
 
     /// `core::ptr::drop_in_place::<T>`.
     ///
@@ -146,7 +146,7 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::FnPtrShim(def_id, _)
             | InstanceDef::Virtual(def_id, _)
             | InstanceDef::Intrinsic(def_id)
-            | InstanceDef::ClosureOnceShim { call_once: def_id }
+            | InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ }
             | InstanceDef::DropGlue(def_id, _)
             | InstanceDef::CloneShim(def_id, _) => def_id,
         }
@@ -161,7 +161,7 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::FnPtrShim(def_id, _)
             | InstanceDef::Virtual(def_id, _)
             | InstanceDef::Intrinsic(def_id)
-            | InstanceDef::ClosureOnceShim { call_once: def_id }
+            | InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ }
             | InstanceDef::DropGlue(def_id, _)
             | InstanceDef::CloneShim(def_id, _) => ty::WithOptConstParam::unknown(def_id),
         }
@@ -231,6 +231,7 @@ impl<'tcx> InstanceDef<'tcx> {
             | InstanceDef::Virtual(def_id, _) => {
                 tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::TRACK_CALLER)
             }
+            InstanceDef::ClosureOnceShim { call_once: _, track_caller } => track_caller,
             _ => false,
         }
     }
@@ -381,6 +382,8 @@ impl<'tcx> Instance<'tcx> {
         substs: SubstsRef<'tcx>,
     ) -> Option<Instance<'tcx>> {
         debug!("resolve(def_id={:?}, substs={:?})", def_id, substs);
+        // Use either `resolve_closure` or `resolve_for_vtable`
+        assert!(!tcx.is_closure(def_id), "Called `resolve_for_fn_ptr` on closure: {:?}", def_id);
         Instance::resolve(tcx, param_env, def_id, substs).ok().flatten().map(|mut resolved| {
             match resolved.def {
                 InstanceDef::Item(def) if resolved.def.requires_caller_location(tcx) => {
@@ -442,12 +445,22 @@ impl<'tcx> Instance<'tcx> {
                                 })
                             )
                         {
+                            if tcx.is_closure(def.did) {
+                                debug!(" => vtable fn pointer created for closure with #[track_caller]: {:?} for method {:?} {:?}",
+                                       def.did, def_id, substs);
+
+                                // Create a shim for the `FnOnce/FnMut/Fn` method we are calling
+                                // - unlike functions, invoking a closure always goes through a
+                                // trait.
+                                resolved = Instance { def: InstanceDef::ReifyShim(def_id), substs };
+                            } else {
                                 debug!(
-                                " => vtable fn pointer created for function with #[track_caller]"
+                                    " => vtable fn pointer created for function with #[track_caller]: {:?}", def.did
                                 );
                                 resolved.def = InstanceDef::ReifyShim(def.did);
                             }
                         }
+                    }
                     InstanceDef::Virtual(def_id, _) => {
                         debug!(" => vtable fn pointer created for virtual call");
                         resolved.def = InstanceDef::ReifyShim(def_id);
@@ -493,7 +506,9 @@ impl<'tcx> Instance<'tcx> {
             .find(|it| it.kind == ty::AssocKind::Fn)
             .unwrap()
             .def_id;
-        let def = ty::InstanceDef::ClosureOnceShim { call_once };
+        let track_caller =
+            tcx.codegen_fn_attrs(closure_did).flags.contains(CodegenFnAttrFlags::TRACK_CALLER);
+        let def = ty::InstanceDef::ClosureOnceShim { call_once, track_caller };
 
         let self_ty = tcx.mk_closure(closure_did, substs);
 

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -638,8 +638,8 @@ impl<'a, 'tcx> Lift<'tcx> for ty::InstanceDef<'a> {
                 Some(ty::InstanceDef::FnPtrShim(def_id, tcx.lift(ty)?))
             }
             ty::InstanceDef::Virtual(def_id, n) => Some(ty::InstanceDef::Virtual(def_id, n)),
-            ty::InstanceDef::ClosureOnceShim { call_once } => {
-                Some(ty::InstanceDef::ClosureOnceShim { call_once })
+            ty::InstanceDef::ClosureOnceShim { call_once, track_caller } => {
+                Some(ty::InstanceDef::ClosureOnceShim { call_once, track_caller })
             }
             ty::InstanceDef::DropGlue(def_id, ty) => {
                 Some(ty::InstanceDef::DropGlue(def_id, tcx.lift(ty)?))
@@ -824,8 +824,8 @@ impl<'tcx> TypeFoldable<'tcx> for ty::instance::Instance<'tcx> {
                 Intrinsic(did) => Intrinsic(did.fold_with(folder)),
                 FnPtrShim(did, ty) => FnPtrShim(did.fold_with(folder), ty.fold_with(folder)),
                 Virtual(did, i) => Virtual(did.fold_with(folder), i),
-                ClosureOnceShim { call_once } => {
-                    ClosureOnceShim { call_once: call_once.fold_with(folder) }
+                ClosureOnceShim { call_once, track_caller } => {
+                    ClosureOnceShim { call_once: call_once.fold_with(folder), track_caller }
                 }
                 DropGlue(did, ty) => DropGlue(did.fold_with(folder), ty.fold_with(folder)),
                 CloneShim(did, ty) => CloneShim(did.fold_with(folder), ty.fold_with(folder)),
@@ -849,7 +849,7 @@ impl<'tcx> TypeFoldable<'tcx> for ty::instance::Instance<'tcx> {
                 did.visit_with(visitor)?;
                 ty.visit_with(visitor)
             }
-            ClosureOnceShim { call_once } => call_once.visit_with(visitor),
+            ClosureOnceShim { call_once, track_caller: _ } => call_once.visit_with(visitor),
         }
     }
 }

compiler/rustc_mir_transform/src/shim.rs

@@ -53,7 +53,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
         ty::InstanceDef::ReifyShim(def_id) => {
             build_call_shim(tcx, instance, None, CallKind::Direct(def_id))
         }
-        ty::InstanceDef::ClosureOnceShim { call_once: _ } => {
+        ty::InstanceDef::ClosureOnceShim { call_once: _, track_caller: _ } => {
             let fn_mut = tcx.require_lang_item(LangItem::FnMut, None);
             let call_mut = tcx
                 .associated_items(fn_mut)

compiler/rustc_span/src/symbol.rs

@@ -407,6 +407,7 @@ symbols! {
         clone_from,
         closure,
         closure_to_fn_coercion,
+        closure_track_caller,
         cmp,
         cmp_max,
         cmp_min,

compiler/rustc_typeck/src/collect.rs

@@ -2778,10 +2778,19 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, id: DefId) -> CodegenFnAttrs {
         } else if attr.has_name(sym::thread_local) {
             codegen_fn_attrs.flags |= CodegenFnAttrFlags::THREAD_LOCAL;
         } else if attr.has_name(sym::track_caller) {
-            if tcx.is_closure(id) || tcx.fn_sig(id).abi() != abi::Abi::Rust {
+            if !tcx.is_closure(id) && tcx.fn_sig(id).abi() != abi::Abi::Rust {
                 struct_span_err!(tcx.sess, attr.span, E0737, "`#[track_caller]` requires Rust ABI")
                     .emit();
             }
+            if tcx.is_closure(id) && !tcx.features().closure_track_caller {
+                feature_err(
+                    &tcx.sess.parse_sess,
+                    sym::closure_track_caller,
+                    attr.span,
+                    "`#[track_caller]` on closures is currently unstable",
+                )
+                .emit();
+            }
             codegen_fn_attrs.flags |= CodegenFnAttrFlags::TRACK_CALLER;
         } else if attr.has_name(sym::export_name) {
             if let Some(s) = attr.value_str() {

src/doc/unstable-book/src/language-features/closure-track-caller.md

@@ -0,0 +1,12 @@
+# `closure_track_caller`
+
+The tracking issue for this feature is: [#87417]
+
+[#87417]: https://github.com/rust-lang/rust/issues/87417
+
+------------------------
+
+Allows using the `#[track_caller]` attribute on closures and generators.
+Calls made to the closure or generator will have caller information
+available through `std::panic::Location::caller()`, just like using
+`#[track_caller]` on a function.

src/test/ui/feature-gates/feature-gate-closure_track_caller.rs

@@ -0,0 +1,7 @@
+#![feature(stmt_expr_attributes)]
+#![feature(generators)]
+
+fn main() {
+    let _closure = #[track_caller] || {}; //~ `#[track_caller]` on closures
+    let _generator = #[track_caller] || { yield; }; //~ `#[track_caller]` on closures
+}

src/test/ui/feature-gates/feature-gate-closure_track_caller.stderr

@@ -0,0 +1,21 @@
+error[E0658]: `#[track_caller]` on closures is currently unstable
+  --> $DIR/feature-gate-closure_track_caller.rs:5:20
+   |
+LL |     let _closure = #[track_caller] || {};
+   |                    ^^^^^^^^^^^^^^^
+   |
+   = note: see issue #87417 <https://github.com/rust-lang/rust/issues/87417> for more information
+   = help: add `#![feature(closure_track_caller)]` to the crate attributes to enable
+
+error[E0658]: `#[track_caller]` on closures is currently unstable
+  --> $DIR/feature-gate-closure_track_caller.rs:6:22
+   |
+LL |     let _generator = #[track_caller] || { yield; };
+   |                      ^^^^^^^^^^^^^^^
+   |
+   = note: see issue #87417 <https://github.com/rust-lang/rust/issues/87417> for more information
+   = help: add `#![feature(closure_track_caller)]` to the crate attributes to enable
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0658`.

src/test/ui/rfc-2091-track-caller/tracked-closure.rs

@@ -0,0 +1,154 @@
+// run-pass
+
+#![feature(stmt_expr_attributes)]
+#![feature(closure_track_caller)]
+#![feature(generator_trait)]
+#![feature(generators)]
+
+use std::ops::{Generator, GeneratorState};
+use std::pin::Pin;
+use std::panic::Location;
+
+type Loc = &'static Location<'static>;
+
+#[track_caller]
+fn mono_invoke_fn<F: Fn(&'static str, bool) -> (&'static str, bool, Loc)>(
+    val: &F
+) -> (&'static str, bool, Loc) {
+    val("from_mono", false)
+}
+
+#[track_caller]
+fn mono_invoke_fn_once<F: FnOnce(&'static str, bool) -> (&'static str, bool, Loc)>(
+    val: F
+) -> (&'static str, bool, Loc) {
+    val("from_mono", false)
+}
+
+#[track_caller]
+fn dyn_invoke_fn_mut(
+    val: &mut dyn FnMut(&'static str, bool) -> (&'static str, bool, Loc)
+) -> (&'static str, bool, Loc) {
+    val("from_dyn", false)
+}
+
+#[track_caller]
+fn dyn_invoke_fn_once(
+    val: Box<dyn FnOnce(&'static str, bool) -> (&'static str, bool, Loc)>
+) -> (&'static str, bool, Loc) {
+    val("from_dyn", false)
+}
+
+
+fn test_closure() {
+    let mut track_closure = #[track_caller] |first: &'static str, second: bool| {
+        (first, second, Location::caller())
+    };
+    let (first_arg, first_bool, first_loc) = track_closure("first_arg", true);
+    let first_line = line!() - 1;
+    assert_eq!(first_arg, "first_arg");
+    assert_eq!(first_bool, true);
+    assert_eq!(first_loc.file(), file!());
+    assert_eq!(first_loc.line(), first_line);
+    assert_eq!(first_loc.column(), 46);
+
+    let (dyn_arg, dyn_bool, dyn_loc) = dyn_invoke_fn_mut(&mut track_closure);
+    assert_eq!(dyn_arg, "from_dyn");
+    assert_eq!(dyn_bool, false);
+    // `FnMut::call_mut` does not have `#[track_caller]`,
+    // so this will not match
+    assert_ne!(dyn_loc.file(), file!());
+
+    let (dyn_arg, dyn_bool, dyn_loc) = dyn_invoke_fn_once(Box::new(track_closure));
+    assert_eq!(dyn_arg, "from_dyn");
+    assert_eq!(dyn_bool, false);
+    // `FnOnce::call_once` does not have `#[track_caller]`
+    // so this will not match
+    assert_ne!(dyn_loc.file(), file!());
+
+
+    let (mono_arg, mono_bool, mono_loc) = mono_invoke_fn(&track_closure);
+    let mono_line = line!() - 1;
+    assert_eq!(mono_arg, "from_mono");
+    assert_eq!(mono_bool, false);
+    assert_eq!(mono_loc.file(), file!());
+    assert_eq!(mono_loc.line(), mono_line);
+    assert_eq!(mono_loc.column(), 43);
+
+    let (mono_arg, mono_bool, mono_loc) = mono_invoke_fn_once(track_closure);
+    let mono_line = line!() - 1;
+    assert_eq!(mono_arg, "from_mono");
+    assert_eq!(mono_bool, false);
+    assert_eq!(mono_loc.file(), file!());
+    assert_eq!(mono_loc.line(), mono_line);
+    assert_eq!(mono_loc.column(), 43);
+
+    let non_tracked_caller = || Location::caller();
+    let non_tracked_line = line!() - 1; // This is the line of the closure, not its caller
+    let non_tracked_loc = non_tracked_caller();
+    assert_eq!(non_tracked_loc.file(), file!());
+    assert_eq!(non_tracked_loc.line(), non_tracked_line);
+    assert_eq!(non_tracked_loc.column(), 33);
+}
+
+
+#[track_caller]
+fn mono_generator<F: Generator<String, Yield = (&'static str, String, Loc), Return = ()>>(
+    val: Pin<&mut F>
+) -> (&'static str, String, Loc) {
+    match val.resume("Mono".to_string()) {
+        GeneratorState::Yielded(val) => val,
+        _ => unreachable!()
+    }
+}
+
+#[track_caller]
+fn dyn_generator(
+    val: Pin<&mut dyn Generator<String, Yield = (&'static str, String, Loc), Return = ()>>
+) -> (&'static str, String, Loc) {
+    match val.resume("Dyn".to_string()) {
+        GeneratorState::Yielded(val) => val,
+        _ => unreachable!()
+    }
+}
+
+fn test_generator() {
+    let generator = #[track_caller] |arg: String| {
+        yield ("first", arg.clone(), Location::caller());
+        yield ("second", arg.clone(), Location::caller());
+    };
+
+    let mut pinned = Box::pin(generator);
+    let (dyn_ret, dyn_arg, dyn_loc) = dyn_generator(pinned.as_mut());
+    assert_eq!(dyn_ret, "first");
+    assert_eq!(dyn_arg, "Dyn".to_string());
+    // The `Generator` trait does not have `#[track_caller]` on `resume`, so
+    // this will not match.
+    assert_ne!(dyn_loc.file(), file!());
+
+
+    let (mono_ret, mono_arg, mono_loc) = mono_generator(pinned.as_mut());
+    let mono_line = line!() - 1;
+    assert_eq!(mono_ret, "second");
+    // The generator ignores the argument to the second `resume` call
+    assert_eq!(mono_arg, "Dyn".to_string());
+    assert_eq!(mono_loc.file(), file!());
+    assert_eq!(mono_loc.line(), mono_line);
+    assert_eq!(mono_loc.column(), 42);
+
+    let non_tracked_generator = || { yield Location::caller(); };
+    let non_tracked_line = line!() - 1; // This is the line of the generator, not its caller
+    let non_tracked_loc = match Box::pin(non_tracked_generator).as_mut().resume(()) {
+        GeneratorState::Yielded(val) => val,
+        _ => unreachable!()
+    };
+    assert_eq!(non_tracked_loc.file(), file!());
+    assert_eq!(non_tracked_loc.line(), non_tracked_line);
+    assert_eq!(non_tracked_loc.column(), 44);
+
+}
+
+fn main() {
+    test_closure();
+    test_generator();
+}