1
Fork 0

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.
This commit is contained in:
Aaron Hill 2021-06-27 14:01:11 -05:00
parent cfff31bc83
commit 94b19fac26
No known key found for this signature in database
GPG key ID: B4087E510E98B164
13 changed files with 267 additions and 23 deletions

View file

@ -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(

View file

@ -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();