Skip to content

Commit 7f02470

Browse files
compiler-errorscompiler-errors
committed
Comments, comments, comments
1 parent cbd1e34 commit 7f02470

File tree

3 files changed

+95
-39
lines changed

3 files changed

+95
-39
lines changed

‎compiler/rustc_mir_transform/src/coroutine/by_move_body.rs

+91-39
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,63 @@
1-
//! A MIR pass which duplicates a coroutine's body and removes any derefs which
2-
//! would be present for upvars that are taken by-ref. The result of which will
3-
//! be a coroutine body that takes all of its upvars by-move, and which we stash
4-
//! into the `CoroutineInfo` for all coroutines returned by coroutine-closures.
1+
//! This pass constructs a second coroutine body sufficient for return from
2+
//! `FnOnce`/`AsyncFnOnce` implementations for coroutine-closures (e.g. async closures).
3+
//!
4+
//! Consider an async closure like:
5+
//! ```rust
6+
//! #![feature(async_closure)]
7+
//!
8+
//! let x = vec![1, 2, 3];
9+
//!
10+
//! let closure = async move || {
11+
//! println!("{x:#?}");
12+
//! };
13+
//! ```
14+
//!
15+
//! This desugars to something like:
16+
//! ```rust,ignore (invalid-borrowck)
17+
//! let x = vec![1, 2, 3];
18+
//!
19+
//! let closure = move || {
20+
//! async {
21+
//! println!("{x:#?}");
22+
//! }
23+
//! };
24+
//! ```
25+
//!
26+
//! Important to note here is that while the outer closure *moves* `x: Vec<i32>`
27+
//! into its upvars, the inner `async` coroutine simply captures a ref of `x`.
28+
//! This is the "magic" of async closures -- the futures that they return are
29+
//! allowed to borrow from their parent closure's upvars.
30+
//!
31+
//! However, what happens when we call `closure` with `AsyncFnOnce` (or `FnOnce`,
32+
//! since all async closures implement that too)? Well, recall the signature:
33+
//! ```
34+
//! pub trait AsyncFnOnce<Args>
35+
//! {
36+
//! type CallOnceFuture: Future<Output = Self::Output>;
37+
//! type Output;
38+
//! fn async_call_once(
39+
//! self,
40+
//! args: Args
41+
//! ) -> Self::CallOnceFuture;
42+
//! }
43+
//! ```
44+
//!
45+
//! This signature *consumes* the async closure (`self`) and returns a `CallOnceFuture`.
46+
//! How do we deal with the fact that the coroutine is supposed to take a reference
47+
//! to the captured `x` from the parent closure, when that parent closure has been
48+
//! destroyed?
49+
//!
50+
//! This is the second piece of magic of async closures. We can simply create a
51+
//! *second* `async` coroutine body where that `x` that was previously captured
52+
//! by reference is now captured by value. This means that we consume the outer
53+
//! closure and return a new coroutine that will hold onto all of these captures,
54+
//! and drop them when it is finished (i.e. after it has been `.await`ed).
55+
//!
56+
//! We do this with the analysis below, which detects the captures that come from
57+
//! borrowing from the outer closure, and we simply peel off a `deref` projection
58+
//! from them. This second body is stored alongside the first body, and optimized
59+
//! with it in lockstep. When we need to resolve a body for `FnOnce` or `AsyncFnOnce`,
60+
//! we use this "by move" body instead.
561
662
use itertools::Itertools;
763

@@ -16,6 +72,8 @@ pub struct ByMoveBody;
1672

1773
impl<'tcx> MirPass<'tcx> for ByMoveBody {
1874
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
75+
// We only need to generate by-move coroutine bodies for coroutines that come
76+
// from coroutine-closures.
1977
let Some(coroutine_def_id) = body.source.def_id().as_local() else {
2078
return;
2179
};
@@ -24,15 +82,19 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
2482
else {
2583
return;
2684
};
85+
86+
// Also, let's skip processing any bodies with errors, since there's no guarantee
87+
// the MIR body will be constructed well.
2788
let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
2889
if coroutine_ty.references_error() {
2990
return;
3091
}
3192

32-
let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
33-
let args = args.as_coroutine();
34-
35-
let coroutine_kind = args.kind_ty().to_opt_closure_kind().unwrap();
93+
let ty::Coroutine(_, coroutine_args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
94+
// We don't need to generate a by-move coroutine if the kind of the coroutine is
95+
// already `FnOnce` -- that means that any upvars that the closure consumes have
96+
// already been taken by-value.
97+
let coroutine_kind = coroutine_args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap();
3698
if coroutine_kind == ty::ClosureKind::FnOnce {
3799
return;
38100
}
@@ -43,12 +105,13 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
43105
else {
44106
bug!();
45107
};
46-
let parent_args = parent_args.as_coroutine_closure();
47-
let parent_upvars_ty = parent_args.tupled_upvars_ty();
48-
let tupled_inputs_ty = tcx.instantiate_bound_regions_with_erased(
49-
parent_args.coroutine_closure_sig().map_bound(|sig| sig.tupled_inputs_ty),
50-
);
51-
let num_args = tupled_inputs_ty.tuple_fields().len();
108+
let parent_closure_args = parent_args.as_coroutine_closure();
109+
let num_args = parent_closure_args
110+
.coroutine_closure_sig()
111+
.skip_binder()
112+
.tupled_inputs_ty
113+
.tuple_fields()
114+
.len();
52115

53116
let mut by_ref_fields = FxIndexSet::default();
54117
for (idx, (coroutine_capture, parent_capture)) in tcx
@@ -59,41 +122,30 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
59122
.zip_eq(tcx.closure_captures(parent_def_id))
60123
.enumerate()
61124
{
62-
// This argument is captured by-move from the parent closure, but by-ref
125+
// This upvar is captured by-move from the parent closure, but by-ref
63126
// from the inner async block. That means that it's being borrowed from
64-
// the closure body -- we need to change the coroutine take it by move.
127+
// the outer closure body -- we need to change the coroutine to take the
128+
// upvar by value.
65129
if coroutine_capture.is_by_ref() && !parent_capture.is_by_ref() {
66130
by_ref_fields.insert(FieldIdx::from_usize(num_args + idx));
67131
}
68132

69133
// Make sure we're actually talking about the same capture.
134+
// FIXME(async_closures): We could look at the `hir::Upvar` instead?
70135
assert_eq!(coroutine_capture.place.ty(), parent_capture.place.ty());
71136
}
72137

73-
let by_move_coroutine_ty = Ty::new_coroutine(
74-
tcx,
75-
coroutine_def_id.to_def_id(),
76-
ty::CoroutineArgs::new(
138+
let by_move_coroutine_ty = tcx
139+
.instantiate_bound_regions_with_erased(parent_closure_args.coroutine_closure_sig())
140+
.to_coroutine_given_kind_and_upvars(
77141
tcx,
78-
ty::CoroutineArgsParts {
79-
parent_args: args.parent_args(),
80-
kind_ty: Ty::from_closure_kind(tcx, ty::ClosureKind::FnOnce),
81-
resume_ty: args.resume_ty(),
82-
yield_ty: args.yield_ty(),
83-
return_ty: args.return_ty(),
84-
witness: args.witness(),
85-
// Concatenate the args + closure's captures (since they're all by move).
86-
tupled_upvars_ty: Ty::new_tup_from_iter(
87-
tcx,
88-
tupled_inputs_ty
89-
.tuple_fields()
90-
.iter()
91-
.chain(parent_upvars_ty.tuple_fields()),
92-
),
93-
},
94-
)
95-
.args,
96-
);
142+
parent_closure_args.parent_args(),
143+
coroutine_def_id.to_def_id(),
144+
ty::ClosureKind::FnOnce,
145+
tcx.lifetimes.re_erased,
146+
parent_closure_args.tupled_upvars_ty(),
147+
parent_closure_args.coroutine_captures_by_ref_ty(),
148+
);
97149

98150
let mut by_move_body = body.clone();
99151
MakeByMoveBody { tcx, by_ref_fields, by_move_coroutine_ty }.visit_body(&mut by_move_body);

‎src/tools/miri/tests/pass/async-closure-captures.rs

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
// Same as rustc's `tests/ui/async-await/async-closures/captures.rs`, keep in sync
2+
13
#![feature(async_closure, noop_waker)]
24

35
use std::future::Future;

‎tests/ui/async-await/async-closures/captures.rs

+2
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
//@ run-pass
44
//@ check-run-results
55

6+
// Same as miri's `tests/pass/async-closure-captures.rs`, keep in sync
7+
68
#![feature(async_closure)]
79

810
extern crate block_on;

0 commit comments

Comments
 (0)