diff options
| author | Michael Goulet <michael@errs.io> | 2024-01-14 22:14:48 +0000 |
|---|---|---|
| committer | Michael Goulet <michael@errs.io> | 2024-01-16 15:15:00 +0000 |
| commit | 8bb1eaee64e170153d21e64cdb437f3d36acd194 (patch) | |
| tree | 5ca6b9e413ec80762443938172ebe348e26c42ee | |
| parent | fa0dc208d0a34027c1d3cca7d47975d8238bcfde (diff) | |
| download | rust-8bb1eaee64e170153d21e64cdb437f3d36acd194.tar.gz rust-8bb1eaee64e170153d21e64cdb437f3d36acd194.zip | |
Introduce helper that deals with moving async args into the coroutine
| -rw-r--r-- | compiler/rustc_ast_lowering/src/item.rs | 363 |
1 files changed, 193 insertions, 170 deletions
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs index a3ff02f5f69..c8e639b7a96 100644 --- a/compiler/rustc_ast_lowering/src/item.rs +++ b/compiler/rustc_ast_lowering/src/item.rs @@ -1082,194 +1082,217 @@ 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(¶meter.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), - }; + let (parameters, expr) = this.lower_coroutine_body_with_moved_arguments( + decl, + body, + coroutine_kind, + CaptureBy::Value { move_kw: rustc_span::DUMMY_SP }, + ); - 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, - ); + // 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); - // 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, - ); + (parameters, expr) + }) + } - statements.push(move_stmt); - statements.push(pattern_stmt); - }; + /// 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); - parameters.push(new_parameter); - } + (ident, false) + } + }; - let mkbody = |this: &mut LoweringContext<'_, 'hir>| { - // Create a block from the user's function body: - let user_body = this.lower_block_expr(body); + let desugared_span = self.mark_span_with_reason(DesugaringKind::Async, span, None); - // 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)); + // 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(¶meter.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), + }; - // As noted above, create the final block like + 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 $param_pattern = $raw_param; - // ... - // drop-temps { <user-body> } - // } + // let __argN = __argN; + // let <pat> = __argN; // ``` - let body = this.block_all( + // + // 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, - this.arena.alloc_from_iter(statements), - Some(user_body), + Some(move_expr), + move_pat, + hir::LocalSource::AsyncFn, ); - 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, + // 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); }; - let coroutine_expr = this.make_desugared_coroutine_expr( - CaptureBy::Value { move_kw: rustc_span::DUMMY_SP }, - closure_id, - None, - body.span, - desugaring_kind, - hir::CoroutineSource::Fn, - mkbody, + + 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), ); - let hir_id = this.lower_node_id(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.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, + ); - (this.arena.alloc_from_iter(parameters), expr) - }) + 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( |