Skip to content

Commit 7c6dc7a

Browse files
davidtwcodavidtwco
committed
Move async fn arguments into closure.
This commit takes advantage of `AsyncArgument` type that was added in a previous commit to replace the arguments of the `async fn` in the HIR and add statements to move the bindings from the new arguments to the pattern from the old argument. For example, the async function `foo` below: async fn foo((x, _y): (T, V)) { async move { } } becomes: async fn foo(__arg0: (T, V)) { async move { let (x, _y) = __arg0; } }
1 parent 879abb1 commit 7c6dc7a

File tree

5 files changed

+115
-37
lines changed

5 files changed

+115
-37
lines changed

‎src/librustc/hir/lowering.rs‎

Lines changed: 67 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -448,10 +448,9 @@ impl<'a> LoweringContext<'a> {
448448
impl<'lcx, 'interner> Visitor<'lcx> for MiscCollector<'lcx, 'interner> {
449449
fn visit_pat(&mut self, p: &'lcx Pat) {
450450
match p.node {
451-
// Doesn't generate a Hir node
451+
// Doesn't generate a HIR node
452452
PatKind::Paren(..) => {},
453453
_ => {
454-
455454
if let Some(owner) = self.hir_id_owner {
456455
self.lctx.lower_node_id_with_owner(p.id, owner);
457456
}
@@ -461,6 +460,31 @@ impl<'a> LoweringContext<'a> {
461460
visit::walk_pat(self, p)
462461
}
463462

463+
fn visit_fn(&mut self, fk: visit::FnKind<'lcx>, fd: &'lcx FnDecl, s: Span, _: NodeId) {
464+
if fk.header().map(|h| h.asyncness.node.is_async()).unwrap_or(false) {
465+
// Don't visit the original pattern for async functions as it will be
466+
// replaced.
467+
for arg in &fd.inputs {
468+
self.visit_ty(&arg.ty)
469+
}
470+
self.visit_fn_ret_ty(&fd.output);
471+
472+
match fk {
473+
visit::FnKind::ItemFn(_, decl, _, body) => {
474+
self.visit_fn_header(decl);
475+
self.visit_block(body)
476+
},
477+
visit::FnKind::Method(_, sig, _, body) => {
478+
self.visit_fn_header(&sig.header);
479+
self.visit_block(body)
480+
},
481+
visit::FnKind::Closure(body) => self.visit_expr(body),
482+
}
483+
} else {
484+
visit::walk_fn(self, fk, fd, s)
485+
}
486+
}
487+
464488
fn visit_item(&mut self, item: &'lcx Item) {
465489
let hir_id = self.lctx.allocate_hir_id_counter(item.id).hir_id;
466490

@@ -3003,12 +3027,18 @@ impl<'a> LoweringContext<'a> {
30033027
asyncness: &IsAsync,
30043028
body: &Block,
30053029
) -> hir::BodyId {
3006-
self.lower_body(Some(decl), |this| {
3007-
if let IsAsync::Async { closure_id, .. } = asyncness {
3030+
self.lower_body(Some(&decl), |this| {
3031+
if let IsAsync::Async { closure_id, ref arguments, .. } = asyncness {
3032+
let mut body = body.clone();
3033+
3034+
for a in arguments.iter().rev() {
3035+
body.stmts.insert(0, a.stmt.clone());
3036+
}
3037+
30083038
let async_expr = this.make_async_expr(
30093039
CaptureBy::Value, *closure_id, None,
30103040
|this| {
3011-
let body = this.lower_block(body, false);
3041+
let body = this.lower_block(&body, false);
30123042
this.expr_block(body, ThinVec::new())
30133043
});
30143044
this.expr(body.span, async_expr, ThinVec::new())
@@ -3070,23 +3100,39 @@ impl<'a> LoweringContext<'a> {
30703100
ItemKind::Fn(ref decl, ref header, ref generics, ref body) => {
30713101
let fn_def_id = self.resolver.definitions().local_def_id(id);
30723102
self.with_new_scopes(|this| {
3073-
// Note: we don't need to change the return type from `T` to
3074-
// `impl Future<Output = T>` here because lower_body
3075-
// only cares about the input argument patterns in the function
3076-
// declaration (decl), not the return types.
3077-
let body_id = this.lower_async_body(decl, &header.asyncness.node, body);
3103+
let mut lower_fn = |decl: &FnDecl| {
3104+
// Note: we don't need to change the return type from `T` to
3105+
// `impl Future<Output = T>` here because lower_body
3106+
// only cares about the input argument patterns in the function
3107+
// declaration (decl), not the return types.
3108+
let body_id = this.lower_async_body(&decl, &header.asyncness.node, body);
3109+
3110+
let (generics, fn_decl) = this.add_in_band_defs(
3111+
generics,
3112+
fn_def_id,
3113+
AnonymousLifetimeMode::PassThrough,
3114+
|this, idty| this.lower_fn_decl(
3115+
&decl,
3116+
Some((fn_def_id, idty)),
3117+
true,
3118+
header.asyncness.node.opt_return_id()
3119+
),
3120+
);
30783121

3079-
let (generics, fn_decl) = this.add_in_band_defs(
3080-
generics,
3081-
fn_def_id,
3082-
AnonymousLifetimeMode::PassThrough,
3083-
|this, idty| this.lower_fn_decl(
3084-
decl,
3085-
Some((fn_def_id, idty)),
3086-
true,
3087-
header.asyncness.node.opt_return_id()
3088-
),
3089-
);
3122+
(body_id, generics, fn_decl)
3123+
};
3124+
3125+
let (body_id, generics, fn_decl) = if let IsAsync::Async {
3126+
arguments, ..
3127+
} = &header.asyncness.node {
3128+
let mut decl = decl.clone();
3129+
// Replace the arguments of this async function with the generated
3130+
// arguments that will be moved into the closure.
3131+
decl.inputs = arguments.clone().drain(..).map(|a| a.arg).collect();
3132+
lower_fn(&decl)
3133+
} else {
3134+
lower_fn(decl)
3135+
};
30903136

30913137
hir::ItemKind::Fn(
30923138
fn_decl,

‎src/librustc/hir/map/def_collector.rs‎

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -73,12 +73,12 @@ impl<'a> DefCollector<'a> {
7373
decl: &'a FnDecl,
7474
body: &'a Block,
7575
) {
76-
let (closure_id, return_impl_trait_id) = match header.asyncness.node {
76+
let (closure_id, return_impl_trait_id, arguments) = match &header.asyncness.node {
7777
IsAsync::Async {
7878
closure_id,
7979
return_impl_trait_id,
80-
..
81-
} => (closure_id, return_impl_trait_id),
80+
arguments,
81+
} => (closure_id, return_impl_trait_id, arguments),
8282
_ => unreachable!(),
8383
};
8484

@@ -87,17 +87,31 @@ impl<'a> DefCollector<'a> {
8787
let fn_def_data = DefPathData::ValueNs(name.as_interned_str());
8888
let fn_def = self.create_def(id, fn_def_data, ITEM_LIKE_SPACE, span);
8989
return self.with_parent(fn_def, |this| {
90-
this.create_def(return_impl_trait_id, DefPathData::ImplTrait, REGULAR_SPACE, span);
90+
this.create_def(*return_impl_trait_id, DefPathData::ImplTrait, REGULAR_SPACE, span);
9191

9292
visit::walk_generics(this, generics);
93-
visit::walk_fn_decl(this, decl);
9493

95-
let closure_def = this.create_def(closure_id,
96-
DefPathData::ClosureExpr,
97-
REGULAR_SPACE,
98-
span);
94+
// Walk the generated arguments for the `async fn`.
95+
for a in arguments {
96+
use visit::Visitor;
97+
this.visit_ty(&a.arg.ty);
98+
}
99+
100+
// We do not invoke `walk_fn_decl` as this will walk the arguments that are being
101+
// replaced.
102+
visit::walk_fn_ret_ty(this, &decl.output);
103+
104+
let closure_def = this.create_def(
105+
*closure_id, DefPathData::ClosureExpr, REGULAR_SPACE, span,
106+
);
99107
this.with_parent(closure_def, |this| {
100-
visit::walk_block(this, body);
108+
for a in arguments {
109+
use visit::Visitor;
110+
// Walk each of the generated statements before the regular block body.
111+
this.visit_stmt(&a.stmt);
112+
}
113+
114+
visit::walk_block(this, &body);
101115
})
102116
})
103117
}

‎src/librustc_resolve/lib.rs‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -834,11 +834,20 @@ impl<'a, 'tcx> Visitor<'tcx> for Resolver<'a> {
834834

835835
// Add each argument to the rib.
836836
let mut bindings_list = FxHashMap::default();
837-
for argument in &declaration.inputs {
837+
let mut add_argument = |argument: &ast::Arg| {
838838
self.resolve_pattern(&argument.pat, PatternSource::FnParam, &mut bindings_list);
839839
self.visit_ty(&argument.ty);
840840
debug!("(resolving function) recorded argument");
841+
};
842+
843+
// Walk the generated async arguments if this is an `async fn`, otherwise walk the
844+
// normal arguments.
845+
if let IsAsync::Async { ref arguments, .. } = asyncness {
846+
for a in arguments { add_argument(&a.arg); }
847+
} else {
848+
for a in &declaration.inputs { add_argument(a); }
841849
}
850+
842851
visit::walk_fn_ret_ty(self, &declaration.output);
843852

844853
// Resolve the function body, potentially inside the body of an async closure
@@ -849,9 +858,18 @@ impl<'a, 'tcx> Visitor<'tcx> for Resolver<'a> {
849858
}
850859

851860
match function_kind {
852-
FnKind::ItemFn(.., body) |
853-
FnKind::Method(.., body) => {
854-
self.visit_block(body);
861+
FnKind::ItemFn(.., body) | FnKind::Method(.., body) => {
862+
if let IsAsync::Async { ref arguments, .. } = asyncness {
863+
let mut body = body.clone();
864+
// Insert the generated statements into the body before attempting to
865+
// resolve names.
866+
for a in arguments {
867+
body.stmts.insert(0, a.stmt.clone());
868+
}
869+
self.visit_block(&body);
870+
} else {
871+
self.visit_block(body);
872+
}
855873
}
856874
FnKind::Closure(body) => {
857875
self.visit_expr(body);

‎src/test/mir-opt/inline-closure-borrows-arg.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ fn foo<T: Copy>(_t: T, q: &i32) -> i32 {
2020
// ...
2121
// bb0: {
2222
// ...
23-
// _3 = [closure@HirId { owner: DefIndex(0:4), local_id: 29 }];
23+
// _3 = [closure@HirId { owner: DefIndex(0:4), local_id: 31 }];
2424
// ...
2525
// _4 = &_3;
2626
// ...

‎src/test/mir-opt/inline-closure.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ fn foo<T: Copy>(_t: T, q: i32) -> i32 {
1616
// ...
1717
// bb0: {
1818
// ...
19-
// _3 = [closure@HirId { owner: DefIndex(0:4), local_id: 13 }];
19+
// _3 = [closure@HirId { owner: DefIndex(0:4), local_id: 15 }];
2020
// ...
2121
// _4 = &_3;
2222
// ...

0 commit comments

Comments
 (0)