dnrops
/
rust
mirror of https://github.com/rust-lang/rust.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Introduce helper that deals with moving async args into the coroutine

This commit is contained in:
Michael Goulet 2024-01-14 22:14:48 +00:00
parent fa0dc208d0
commit 8bb1eaee64
1 changed files with 204 additions and 181 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

385
compiler/rustc_ast_lowering/src/item.rs
View File

@ -1082,196 +1082,219 @@ impl<'hir> LoweringContext<'_, 'hir> {
let (Some(coroutine_kind), Some(body)) = (coroutine_kind, body) else {
return self.lower_fn_body_block(span, decl, body);
};
let closure_id = coroutine_kind.closure_id();
self.lower_body(|this| {
let mut parameters: Vec<hir::Param<'_>> = Vec::new();
let mut statements: Vec<hir::Stmt<'_>> = Vec::new();
// Async function parameters are lowered into the closure body so that they are
// captured and so that the drop order matches the equivalent non-async functions.
//
// from:
//
// async fn foo(<pattern>: <ty>, <pattern>: <ty>, <pattern>: <ty>) {
// <body>
// }
//
// into:
//
// fn foo(__arg0: <ty>, __arg1: <ty>, __arg2: <ty>) {
// async move {
// let __arg2 = __arg2;
// let <pattern> = __arg2;
// let __arg1 = __arg1;
// let <pattern> = __arg1;
// let __arg0 = __arg0;
// let <pattern> = __arg0;
// drop-temps { <body> } // see comments later in fn for details
// }
// }
//
// If `<pattern>` is a simple ident, then it is lowered to a single
// `let <pattern> = <pattern>;` statement as an optimization.
//
// Note that the body is embedded in `drop-temps`; an
// equivalent desugaring would be `return { <body>
// };`. The key point is that we wish to drop all the
// let-bound variables and temporaries created in the body
// (and its tail expression!) before we drop the
// parameters (c.f. rust-lang/rust#64512).
for (index, parameter) in decl.inputs.iter().enumerate() {
let parameter = this.lower_param(parameter);
let span = parameter.pat.span;
// Check if this is a binding pattern, if so, we can optimize and avoid adding a
// `let <pat> = __argN;` statement. In this case, we do not rename the parameter.
let (ident, is_simple_parameter) = match parameter.pat.kind {
hir::PatKind::Binding(hir::BindingAnnotation(ByRef::No, _), _, ident, _) => {
(ident, true)
}
// For `ref mut` or wildcard arguments, we can't reuse the binding, but
// we can keep the same name for the parameter.
// This lets rustdoc render it correctly in documentation.
hir::PatKind::Binding(_, _, ident, _) => (ident, false),
hir::PatKind::Wild => {
(Ident::with_dummy_span(rustc_span::symbol::kw::Underscore), false)
}
_ => {
// Replace the ident for bindings that aren't simple.
let name = format!("__arg{index}");
let ident = Ident::from_str(&name);
(ident, false)
}
};
let desugared_span = this.mark_span_with_reason(DesugaringKind::Async, span, None);
// Construct a parameter representing `__argN: <ty>` to replace the parameter of the
// async function.
//
// If this is the simple case, this parameter will end up being the same as the
// original parameter, but with a different pattern id.
let stmt_attrs = this.attrs.get(&parameter.hir_id.local_id).copied();
let (new_parameter_pat, new_parameter_id) = this.pat_ident(desugared_span, ident);
let new_parameter = hir::Param {
hir_id: parameter.hir_id,
pat: new_parameter_pat,
ty_span: this.lower_span(parameter.ty_span),
span: this.lower_span(parameter.span),
};
if is_simple_parameter {
// If this is the simple case, then we only insert one statement that is
// `let <pat> = <pat>;`. We re-use the original argument's pattern so that
// `HirId`s are densely assigned.
let expr = this.expr_ident(desugared_span, ident, new_parameter_id);
let stmt = this.stmt_let_pat(
stmt_attrs,
desugared_span,
Some(expr),
parameter.pat,
hir::LocalSource::AsyncFn,
);
statements.push(stmt);
} else {
// If this is not the simple case, then we construct two statements:
//
// ```
// let __argN = __argN;
// let <pat> = __argN;
// ```
//
// The first statement moves the parameter into the closure and thus ensures
// that the drop order is correct.
//
// The second statement creates the bindings that the user wrote.
// Construct the `let mut __argN = __argN;` statement. It must be a mut binding
// because the user may have specified a `ref mut` binding in the next
// statement.
let (move_pat, move_id) = this.pat_ident_binding_mode(
desugared_span,
ident,
hir::BindingAnnotation::MUT,
);
let move_expr = this.expr_ident(desugared_span, ident, new_parameter_id);
let move_stmt = this.stmt_let_pat(
None,
desugared_span,
Some(move_expr),
move_pat,
hir::LocalSource::AsyncFn,
);
// Construct the `let <pat> = __argN;` statement. We re-use the original
// parameter's pattern so that `HirId`s are densely assigned.
let pattern_expr = this.expr_ident(desugared_span, ident, move_id);
let pattern_stmt = this.stmt_let_pat(
stmt_attrs,
desugared_span,
Some(pattern_expr),
parameter.pat,
hir::LocalSource::AsyncFn,
);
statements.push(move_stmt);
statements.push(pattern_stmt);
};
parameters.push(new_parameter);
}
let mkbody = |this: &mut LoweringContext<'_, 'hir>| {
// Create a block from the user's function body:
let user_body = this.lower_block_expr(body);
// Transform into `drop-temps { <user-body> }`, an expression:
let desugared_span =
this.mark_span_with_reason(DesugaringKind::Async, user_body.span, None);
let user_body = this.expr_drop_temps(desugared_span, this.arena.alloc(user_body));
// As noted above, create the final block like
//
// ```
// {
// let $param_pattern = $raw_param;
// ...
// drop-temps { <user-body> }
// }
// ```
let body = this.block_all(
desugared_span,
this.arena.alloc_from_iter(statements),
Some(user_body),
);
this.expr_block(body)
};
let desugaring_kind = match coroutine_kind {
CoroutineKind::Async { .. } => hir::CoroutineDesugaring::Async,
CoroutineKind::Gen { .. } => hir::CoroutineDesugaring::Gen,
CoroutineKind::AsyncGen { .. } => hir::CoroutineDesugaring::AsyncGen,
};
let coroutine_expr = this.make_desugared_coroutine_expr(
let (parameters, expr) = this.lower_coroutine_body_with_moved_arguments(
decl,
body,
coroutine_kind,
CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
closure_id,
None,
body.span,
desugaring_kind,
hir::CoroutineSource::Fn,
mkbody,
);
let hir_id = this.lower_node_id(closure_id);
// FIXME(async_fn_track_caller): Can this be moved above?
let hir_id = this.lower_node_id(coroutine_kind.closure_id());
this.maybe_forward_track_caller(body.span, fn_id, hir_id);
let expr = hir::Expr { hir_id, kind: coroutine_expr, span: this.lower_span(body.span) };
(this.arena.alloc_from_iter(parameters), expr)
(parameters, expr)
})
}
/// Lowers a desugared coroutine body after moving all of the arguments
/// into the body. This is to make sure that the future actually owns the
/// arguments that are passed to the function, and to ensure things like
/// drop order are stable.
fn lower_coroutine_body_with_moved_arguments(
&mut self,
decl: &FnDecl,
body: &Block,
coroutine_kind: CoroutineKind,
capture_clause: CaptureBy,
) -> (&'hir [hir::Param<'hir>], hir::Expr<'hir>) {
let mut parameters: Vec<hir::Param<'_>> = Vec::new();
let mut statements: Vec<hir::Stmt<'_>> = Vec::new();
// Async function parameters are lowered into the closure body so that they are
// captured and so that the drop order matches the equivalent non-async functions.
//
// from:
//
// async fn foo(<pattern>: <ty>, <pattern>: <ty>, <pattern>: <ty>) {
// <body>
// }
//
// into:
//
// fn foo(__arg0: <ty>, __arg1: <ty>, __arg2: <ty>) {
// async move {
// let __arg2 = __arg2;
// let <pattern> = __arg2;
// let __arg1 = __arg1;
// let <pattern> = __arg1;
// let __arg0 = __arg0;
// let <pattern> = __arg0;
// drop-temps { <body> } // see comments later in fn for details
// }
// }
//
// If `<pattern>` is a simple ident, then it is lowered to a single
// `let <pattern> = <pattern>;` statement as an optimization.
//
// Note that the body is embedded in `drop-temps`; an
// equivalent desugaring would be `return { <body>
// };`. The key point is that we wish to drop all the
// let-bound variables and temporaries created in the body
// (and its tail expression!) before we drop the
// parameters (c.f. rust-lang/rust#64512).
for (index, parameter) in decl.inputs.iter().enumerate() {
let parameter = self.lower_param(parameter);
let span = parameter.pat.span;
// Check if this is a binding pattern, if so, we can optimize and avoid adding a
// `let <pat> = __argN;` statement. In this case, we do not rename the parameter.
let (ident, is_simple_parameter) = match parameter.pat.kind {
hir::PatKind::Binding(hir::BindingAnnotation(ByRef::No, _), _, ident, _) => {
(ident, true)
}
// For `ref mut` or wildcard arguments, we can't reuse the binding, but
// we can keep the same name for the parameter.
// This lets rustdoc render it correctly in documentation.
hir::PatKind::Binding(_, _, ident, _) => (ident, false),
hir::PatKind::Wild => {
(Ident::with_dummy_span(rustc_span::symbol::kw::Underscore), false)
}
_ => {
// Replace the ident for bindings that aren't simple.
let name = format!("__arg{index}");
let ident = Ident::from_str(&name);
(ident, false)
}
};
let desugared_span = self.mark_span_with_reason(DesugaringKind::Async, span, None);
// Construct a parameter representing `__argN: <ty>` to replace the parameter of the
// async function.
//
// If this is the simple case, this parameter will end up being the same as the
// original parameter, but with a different pattern id.
let stmt_attrs = self.attrs.get(&parameter.hir_id.local_id).copied();
let (new_parameter_pat, new_parameter_id) = self.pat_ident(desugared_span, ident);
let new_parameter = hir::Param {
hir_id: parameter.hir_id,
pat: new_parameter_pat,
ty_span: self.lower_span(parameter.ty_span),
span: self.lower_span(parameter.span),
};
if is_simple_parameter {
// If this is the simple case, then we only insert one statement that is
// `let <pat> = <pat>;`. We re-use the original argument's pattern so that
// `HirId`s are densely assigned.
let expr = self.expr_ident(desugared_span, ident, new_parameter_id);
let stmt = self.stmt_let_pat(
stmt_attrs,
desugared_span,
Some(expr),
parameter.pat,
hir::LocalSource::AsyncFn,
);
statements.push(stmt);
} else {
// If this is not the simple case, then we construct two statements:
//
// ```
// let __argN = __argN;
// let <pat> = __argN;
// ```
//
// The first statement moves the parameter into the closure and thus ensures
// that the drop order is correct.
//
// The second statement creates the bindings that the user wrote.
// Construct the `let mut __argN = __argN;` statement. It must be a mut binding
// because the user may have specified a `ref mut` binding in the next
// statement.
let (move_pat, move_id) =
self.pat_ident_binding_mode(desugared_span, ident, hir::BindingAnnotation::MUT);
let move_expr = self.expr_ident(desugared_span, ident, new_parameter_id);
let move_stmt = self.stmt_let_pat(
None,
desugared_span,
Some(move_expr),
move_pat,
hir::LocalSource::AsyncFn,
);
// Construct the `let <pat> = __argN;` statement. We re-use the original
// parameter's pattern so that `HirId`s are densely assigned.
let pattern_expr = self.expr_ident(desugared_span, ident, move_id);
let pattern_stmt = self.stmt_let_pat(
stmt_attrs,
desugared_span,
Some(pattern_expr),
parameter.pat,
hir::LocalSource::AsyncFn,
);
statements.push(move_stmt);
statements.push(pattern_stmt);
};
parameters.push(new_parameter);
}
let mkbody = |this: &mut LoweringContext<'_, 'hir>| {
// Create a block from the user's function body:
let user_body = this.lower_block_expr(body);
// Transform into `drop-temps { <user-body> }`, an expression:
let desugared_span =
this.mark_span_with_reason(DesugaringKind::Async, user_body.span, None);
let user_body = this.expr_drop_temps(desugared_span, this.arena.alloc(user_body));
// As noted above, create the final block like
//
// ```
// {
// let $param_pattern = $raw_param;
// ...
// drop-temps { <user-body> }
// }
// ```
let body = this.block_all(
desugared_span,
this.arena.alloc_from_iter(statements),
Some(user_body),
);
this.expr_block(body)
};
let desugaring_kind = match coroutine_kind {
CoroutineKind::Async { .. } => hir::CoroutineDesugaring::Async,
CoroutineKind::Gen { .. } => hir::CoroutineDesugaring::Gen,
CoroutineKind::AsyncGen { .. } => hir::CoroutineDesugaring::AsyncGen,
};
let closure_id = coroutine_kind.closure_id();
let coroutine_expr = self.make_desugared_coroutine_expr(
capture_clause,
closure_id,
None,
body.span,
desugaring_kind,
hir::CoroutineSource::Fn,
mkbody,
);
let expr = hir::Expr {
hir_id: self.lower_node_id(closure_id),
kind: coroutine_expr,
span: self.lower_span(body.span),
};
(self.arena.alloc_from_iter(parameters), expr)
}
fn lower_method_sig(
&mut self,
generics: &Generics,