Skip to content

Commit 0132f82

Browse files
borsbors
committed
Auto merge of #87064 - Aaron1011:new-closure-track-caller, r=estebank
Support `#[track_caller]` on closures and generators ## Lang team summary This PR adds support for placing the `#[track_caller]` attribute on closure and generator expressions. This attribute's addition behaves identically (from a users perspective) to the attribute being placed on the method in impl Fn/FnOnce/FnMut for ... generated by compiler. The attribute is currently "double" feature gated -- both `stmt_expr_attributes` (preexisting) and `closure_track_caller` (newly added) must be enabled in order to place these attributes on closures. As the Fn* traits lack a `#[track_caller]` attribute in their definition, caller information does not propagate when invoking closures through dyn Fn*. There is no limitation that this PR adds in supporting this; it can be added in the future. # Implementation details 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.
2 parents 15d9ba0 + 94b19fa commit 0132f82

File tree

13 files changed

+267
-23
lines changed

13 files changed

+267
-23
lines changed

‎compiler/rustc_codegen_ssa/src/mir/block.rs

+14-5
Original file line numberDiff line numberDiff line change
@@ -777,22 +777,30 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
777777

778778
self.codegen_argument(&mut bx, op, &mut llargs, &fn_abi.args[i]);
779779
}
780-
if let Some(tup) = untuple {
780+
let num_untupled = untuple.map(|tup| {
781781
self.codegen_arguments_untupled(
782782
&mut bx,
783783
tup,
784784
&mut llargs,
785785
&fn_abi.args[first_args.len()..],
786786
)
787-
}
787+
});
788788

789789
let needs_location =
790790
instance.map_or(false, |i| i.def.requires_caller_location(self.cx.tcx()));
791791
if needs_location {
792+
let mir_args = if let Some(num_untupled) = num_untupled {
793+
first_args.len() + num_untupled
794+
} else {
795+
args.len()
796+
};
792797
assert_eq!(
793798
fn_abi.args.len(),
794-
args.len() + 1,
795-
"#[track_caller] fn's must have 1 more argument in their ABI than in their MIR",
799+
mir_args + 1,
800+
"#[track_caller] fn's must have 1 more argument in their ABI than in their MIR: {:?} {:?} {:?}",
801+
instance,
802+
fn_span,
803+
fn_abi,
796804
);
797805
let location =
798806
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> {
11221130
operand: &mir::Operand<'tcx>,
11231131
llargs: &mut Vec<Bx::Value>,
11241132
args: &[ArgAbi<'tcx, Ty<'tcx>>],
1125-
) {
1133+
) -> usize {
11261134
let tuple = self.codegen_operand(bx, operand);
11271135

11281136
// Handle both by-ref and immediate tuples.
@@ -1142,6 +1150,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
11421150
self.codegen_argument(bx, op, llargs, &args[i]);
11431151
}
11441152
}
1153+
tuple.layout.fields.count()
11451154
}
11461155

11471156
fn get_caller_location(

‎compiler/rustc_codegen_ssa/src/mir/mod.rs

+16-2
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,8 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
258258
let mut idx = 0;
259259
let mut llarg_idx = fx.fn_abi.ret.is_indirect() as usize;
260260

261+
let mut num_untupled = None;
262+
261263
let args = mir
262264
.args_iter()
263265
.enumerate()
@@ -286,6 +288,11 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
286288
let pr_field = place.project_field(bx, i);
287289
bx.store_fn_arg(arg, &mut llarg_idx, pr_field);
288290
}
291+
assert_eq!(
292+
None,
293+
num_untupled.replace(tupled_arg_tys.len()),
294+
"Replaced existing num_tupled"
295+
);
289296

290297
return LocalRef::Place(place);
291298
}
@@ -362,10 +369,17 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
362369
.collect::<Vec<_>>();
363370

364371
if fx.instance.def.requires_caller_location(bx.tcx()) {
372+
let mir_args = if let Some(num_untupled) = num_untupled {
373+
// Subtract off the tupled argument that gets 'expanded'
374+
args.len() - 1 + num_untupled
375+
} else {
376+
args.len()
377+
};
365378
assert_eq!(
366379
fx.fn_abi.args.len(),
367-
args.len() + 1,
368-
"#[track_caller] fn's must have 1 more argument in their ABI than in their MIR",
380+
mir_args + 1,
381+
"#[track_caller] instance {:?} must have 1 more argument in their ABI than in their MIR",
382+
fx.instance
369383
);
370384

371385
let arg = fx.fn_abi.args.last().unwrap();

‎compiler/rustc_feature/src/active.rs

+2
Original file line numberDiff line numberDiff line change
@@ -678,6 +678,8 @@ declare_features! (
678678
/// Allows the `#[must_not_suspend]` attribute.
679679
(active, must_not_suspend, "1.57.0", Some(83310), None),
680680

681+
/// Allows `#[track_caller]` on closures and generators.
682+
(active, closure_track_caller, "1.57.0", Some(87417), None),
681683

682684
// -------------------------------------------------------------------------
683685
// feature-group-end: actual feature gates

‎compiler/rustc_middle/src/mir/visit.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -348,7 +348,7 @@ macro_rules! make_mir_visitor {
348348
ty::InstanceDef::VtableShim(_def_id) |
349349
ty::InstanceDef::ReifyShim(_def_id) |
350350
ty::InstanceDef::Virtual(_def_id, _) |
351-
ty::InstanceDef::ClosureOnceShim { call_once: _def_id } |
351+
ty::InstanceDef::ClosureOnceShim { call_once: _def_id, track_caller: _ } |
352352
ty::InstanceDef::DropGlue(_def_id, None) => {}
353353

354354
ty::InstanceDef::FnPtrShim(_def_id, ty) |

‎compiler/rustc_middle/src/ty/instance.rs

+23-8
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ pub enum InstanceDef<'tcx> {
7777
/// `<[FnMut closure] as FnOnce>::call_once`.
7878
///
7979
/// The `DefId` is the ID of the `call_once` method in `FnOnce`.
80-
ClosureOnceShim { call_once: DefId },
80+
ClosureOnceShim { call_once: DefId, track_caller: bool },
8181

8282
/// `core::ptr::drop_in_place::<T>`.
8383
///
@@ -146,7 +146,7 @@ impl<'tcx> InstanceDef<'tcx> {
146146
| InstanceDef::FnPtrShim(def_id, _)
147147
| InstanceDef::Virtual(def_id, _)
148148
| InstanceDef::Intrinsic(def_id)
149-
| InstanceDef::ClosureOnceShim { call_once: def_id }
149+
| InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ }
150150
| InstanceDef::DropGlue(def_id, _)
151151
| InstanceDef::CloneShim(def_id, _) => def_id,
152152
}
@@ -161,7 +161,7 @@ impl<'tcx> InstanceDef<'tcx> {
161161
| InstanceDef::FnPtrShim(def_id, _)
162162
| InstanceDef::Virtual(def_id, _)
163163
| InstanceDef::Intrinsic(def_id)
164-
| InstanceDef::ClosureOnceShim { call_once: def_id }
164+
| InstanceDef::ClosureOnceShim { call_once: def_id, track_caller: _ }
165165
| InstanceDef::DropGlue(def_id, _)
166166
| InstanceDef::CloneShim(def_id, _) => ty::WithOptConstParam::unknown(def_id),
167167
}
@@ -231,6 +231,7 @@ impl<'tcx> InstanceDef<'tcx> {
231231
| InstanceDef::Virtual(def_id, _) => {
232232
tcx.codegen_fn_attrs(def_id).flags.contains(CodegenFnAttrFlags::TRACK_CALLER)
233233
}
234+
InstanceDef::ClosureOnceShim { call_once: _, track_caller } => track_caller,
234235
_ => false,
235236
}
236237
}
@@ -381,6 +382,8 @@ impl<'tcx> Instance<'tcx> {
381382
substs: SubstsRef<'tcx>,
382383
) -> Option<Instance<'tcx>> {
383384
debug!("resolve(def_id={:?}, substs={:?})", def_id, substs);
385+
// Use either `resolve_closure` or `resolve_for_vtable`
386+
assert!(!tcx.is_closure(def_id), "Called `resolve_for_fn_ptr` on closure: {:?}", def_id);
384387
Instance::resolve(tcx, param_env, def_id, substs).ok().flatten().map(|mut resolved| {
385388
match resolved.def {
386389
InstanceDef::Item(def) if resolved.def.requires_caller_location(tcx) => {
@@ -442,10 +445,20 @@ impl<'tcx> Instance<'tcx> {
442445
})
443446
)
444447
{
445-
debug!(
446-
" => vtable fn pointer created for function with #[track_caller]"
447-
);
448-
resolved.def = InstanceDef::ReifyShim(def.did);
448+
if tcx.is_closure(def.did) {
449+
debug!(" => vtable fn pointer created for closure with #[track_caller]: {:?} for method {:?} {:?}",
450+
def.did, def_id, substs);
451+
452+
// Create a shim for the `FnOnce/FnMut/Fn` method we are calling
453+
// - unlike functions, invoking a closure always goes through a
454+
// trait.
455+
resolved = Instance { def: InstanceDef::ReifyShim(def_id), substs };
456+
} else {
457+
debug!(
458+
" => vtable fn pointer created for function with #[track_caller]: {:?}", def.did
459+
);
460+
resolved.def = InstanceDef::ReifyShim(def.did);
461+
}
449462
}
450463
}
451464
InstanceDef::Virtual(def_id, _) => {
@@ -493,7 +506,9 @@ impl<'tcx> Instance<'tcx> {
493506
.find(|it| it.kind == ty::AssocKind::Fn)
494507
.unwrap()
495508
.def_id;
496-
let def = ty::InstanceDef::ClosureOnceShim { call_once };
509+
let track_caller =
510+
tcx.codegen_fn_attrs(closure_did).flags.contains(CodegenFnAttrFlags::TRACK_CALLER);
511+
let def = ty::InstanceDef::ClosureOnceShim { call_once, track_caller };
497512

498513
let self_ty = tcx.mk_closure(closure_did, substs);
499514

‎compiler/rustc_middle/src/ty/structural_impls.rs

+5-5
Original file line numberDiff line numberDiff line change
@@ -638,8 +638,8 @@ impl<'a, 'tcx> Lift<'tcx> for ty::InstanceDef<'a> {
638638
Some(ty::InstanceDef::FnPtrShim(def_id, tcx.lift(ty)?))
639639
}
640640
ty::InstanceDef::Virtual(def_id, n) => Some(ty::InstanceDef::Virtual(def_id, n)),
641-
ty::InstanceDef::ClosureOnceShim { call_once } => {
642-
Some(ty::InstanceDef::ClosureOnceShim { call_once })
641+
ty::InstanceDef::ClosureOnceShim { call_once, track_caller } => {
642+
Some(ty::InstanceDef::ClosureOnceShim { call_once, track_caller })
643643
}
644644
ty::InstanceDef::DropGlue(def_id, ty) => {
645645
Some(ty::InstanceDef::DropGlue(def_id, tcx.lift(ty)?))
@@ -824,8 +824,8 @@ impl<'tcx> TypeFoldable<'tcx> for ty::instance::Instance<'tcx> {
824824
Intrinsic(did) => Intrinsic(did.fold_with(folder)),
825825
FnPtrShim(did, ty) => FnPtrShim(did.fold_with(folder), ty.fold_with(folder)),
826826
Virtual(did, i) => Virtual(did.fold_with(folder), i),
827-
ClosureOnceShim { call_once } => {
828-
ClosureOnceShim { call_once: call_once.fold_with(folder) }
827+
ClosureOnceShim { call_once, track_caller } => {
828+
ClosureOnceShim { call_once: call_once.fold_with(folder), track_caller }
829829
}
830830
DropGlue(did, ty) => DropGlue(did.fold_with(folder), ty.fold_with(folder)),
831831
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> {
849849
did.visit_with(visitor)?;
850850
ty.visit_with(visitor)
851851
}
852-
ClosureOnceShim { call_once } => call_once.visit_with(visitor),
852+
ClosureOnceShim { call_once, track_caller: _ } => call_once.visit_with(visitor),
853853
}
854854
}
855855
}

‎compiler/rustc_mir_transform/src/shim.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
5353
ty::InstanceDef::ReifyShim(def_id) => {
5454
build_call_shim(tcx, instance, None, CallKind::Direct(def_id))
5555
}
56-
ty::InstanceDef::ClosureOnceShim { call_once: _ } => {
56+
ty::InstanceDef::ClosureOnceShim { call_once: _, track_caller: _ } => {
5757
let fn_mut = tcx.require_lang_item(LangItem::FnMut, None);
5858
let call_mut = tcx
5959
.associated_items(fn_mut)

‎compiler/rustc_span/src/symbol.rs

+1
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,7 @@ symbols! {
408408
clone_from,
409409
closure,
410410
closure_to_fn_coercion,
411+
closure_track_caller,
411412
cmp,
412413
cmp_max,
413414
cmp_min,

‎compiler/rustc_typeck/src/collect.rs

+10-1
Original file line numberDiff line numberDiff line change
@@ -2778,10 +2778,19 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, id: DefId) -> CodegenFnAttrs {
27782778
} else if attr.has_name(sym::thread_local) {
27792779
codegen_fn_attrs.flags |= CodegenFnAttrFlags::THREAD_LOCAL;
27802780
} else if attr.has_name(sym::track_caller) {
2781-
if tcx.is_closure(id) || tcx.fn_sig(id).abi() != abi::Abi::Rust {
2781+
if !tcx.is_closure(id) && tcx.fn_sig(id).abi() != abi::Abi::Rust {
27822782
struct_span_err!(tcx.sess, attr.span, E0737, "`#[track_caller]` requires Rust ABI")
27832783
.emit();
27842784
}
2785+
if tcx.is_closure(id) && !tcx.features().closure_track_caller {
2786+
feature_err(
2787+
&tcx.sess.parse_sess,
2788+
sym::closure_track_caller,
2789+
attr.span,
2790+
"`#[track_caller]` on closures is currently unstable",
2791+
)
2792+
.emit();
2793+
}
27852794
codegen_fn_attrs.flags |= CodegenFnAttrFlags::TRACK_CALLER;
27862795
} else if attr.has_name(sym::export_name) {
27872796
if let Some(s) = attr.value_str() {

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

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
# `closure_track_caller`
2+
3+
The tracking issue for this feature is: [#87417]
4+
5+
[#87417]: https://github.com/rust-lang/rust/issues/87417
6+
7+
------------------------
8+
9+
Allows using the `#[track_caller]` attribute on closures and generators.
10+
Calls made to the closure or generator will have caller information
11+
available through `std::panic::Location::caller()`, just like using
12+
`#[track_caller]` on a function.

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

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
#![feature(stmt_expr_attributes)]
2+
#![feature(generators)]
3+
4+
fn main() {
5+
let _closure = #[track_caller] || {}; //~ `#[track_caller]` on closures
6+
let _generator = #[track_caller] || { yield; }; //~ `#[track_caller]` on closures
7+
}

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

+21
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
error[E0658]: `#[track_caller]` on closures is currently unstable
2+
--> $DIR/feature-gate-closure_track_caller.rs:5:20
3+
|
4+
LL | let _closure = #[track_caller] || {};
5+
| ^^^^^^^^^^^^^^^
6+
|
7+
= note: see issue #87417 <https://github.com/rust-lang/rust/issues/87417> for more information
8+
= help: add `#![feature(closure_track_caller)]` to the crate attributes to enable
9+
10+
error[E0658]: `#[track_caller]` on closures is currently unstable
11+
--> $DIR/feature-gate-closure_track_caller.rs:6:22
12+
|
13+
LL | let _generator = #[track_caller] || { yield; };
14+
| ^^^^^^^^^^^^^^^
15+
|
16+
= note: see issue #87417 <https://github.com/rust-lang/rust/issues/87417> for more information
17+
= help: add `#![feature(closure_track_caller)]` to the crate attributes to enable
18+
19+
error: aborting due to 2 previous errors
20+
21+
For more information about this error, try `rustc --explain E0658`.

0 commit comments

Comments
 (0)