Skip to content

Commit 4b4bcf4

Browse files
Introduce helper that deals with moving async args into the coroutine
1 parent 5876c8c commit 4b4bcf4

File tree

1 file changed

+191
-170
lines changed
  • compiler/rustc_ast_lowering/src

1 file changed

+191
-170
lines changed

compiler/rustc_ast_lowering/src/item.rs

Lines changed: 191 additions & 170 deletions
Original file line numberDiff line numberDiff line change
@@ -1049,194 +1049,215 @@ impl<'hir> LoweringContext<'_, 'hir> {
10491049
let (Some(coroutine_kind), Some(body)) = (coroutine_kind, body) else {
10501050
return self.lower_fn_body_block(span, decl, body);
10511051
};
1052-
let closure_id = coroutine_kind.closure_id();
1053-
10541052
self.lower_body(|this| {
1055-
let mut parameters: Vec<hir::Param<'_>> = Vec::new();
1056-
let mut statements: Vec<hir::Stmt<'_>> = Vec::new();
1057-
1058-
// Async function parameters are lowered into the closure body so that they are
1059-
// captured and so that the drop order matches the equivalent non-async functions.
1060-
//
1061-
// from:
1062-
//
1063-
// async fn foo(<pattern>: <ty>, <pattern>: <ty>, <pattern>: <ty>) {
1064-
// <body>
1065-
// }
1066-
//
1067-
// into:
1068-
//
1069-
// fn foo(__arg0: <ty>, __arg1: <ty>, __arg2: <ty>) {
1070-
// async move {
1071-
// let __arg2 = __arg2;
1072-
// let <pattern> = __arg2;
1073-
// let __arg1 = __arg1;
1074-
// let <pattern> = __arg1;
1075-
// let __arg0 = __arg0;
1076-
// let <pattern> = __arg0;
1077-
// drop-temps { <body> } // see comments later in fn for details
1078-
// }
1079-
// }
1080-
//
1081-
// If `<pattern>` is a simple ident, then it is lowered to a single
1082-
// `let <pattern> = <pattern>;` statement as an optimization.
1083-
//
1084-
// Note that the body is embedded in `drop-temps`; an
1085-
// equivalent desugaring would be `return { <body>
1086-
// };`. The key point is that we wish to drop all the
1087-
// let-bound variables and temporaries created in the body
1088-
// (and its tail expression!) before we drop the
1089-
// parameters (c.f. rust-lang/rust#64512).
1090-
for (index, parameter) in decl.inputs.iter().enumerate() {
1091-
let parameter = this.lower_param(parameter);
1092-
let span = parameter.pat.span;
1093-
1094-
// Check if this is a binding pattern, if so, we can optimize and avoid adding a
1095-
// `let <pat> = __argN;` statement. In this case, we do not rename the parameter.
1096-
let (ident, is_simple_parameter) = match parameter.pat.kind {
1097-
hir::PatKind::Binding(hir::BindingAnnotation(ByRef::No, _), _, ident, _) => {
1098-
(ident, true)
1099-
}
1100-
// For `ref mut` or wildcard arguments, we can't reuse the binding, but
1101-
// we can keep the same name for the parameter.
1102-
// This lets rustdoc render it correctly in documentation.
1103-
hir::PatKind::Binding(_, _, ident, _) => (ident, false),
1104-
hir::PatKind::Wild => {
1105-
(Ident::with_dummy_span(rustc_span::symbol::kw::Underscore), false)
1106-
}
1107-
_ => {
1108-
// Replace the ident for bindings that aren't simple.
1109-
let name = format!("__arg{index}");
1110-
let ident = Ident::from_str(&name);
1111-
1112-
(ident, false)
1113-
}
1114-
};
1115-
1116-
let desugared_span = this.mark_span_with_reason(DesugaringKind::Async, span, None);
1117-
1118-
// Construct a parameter representing `__argN: <ty>` to replace the parameter of the
1119-
// async function.
1120-
//
1121-
// If this is the simple case, this parameter will end up being the same as the
1122-
// original parameter, but with a different pattern id.
1123-
let stmt_attrs = this.attrs.get(&parameter.hir_id.local_id).copied();
1124-
let (new_parameter_pat, new_parameter_id) = this.pat_ident(desugared_span, ident);
1125-
let new_parameter = hir::Param {
1126-
hir_id: parameter.hir_id,
1127-
pat: new_parameter_pat,
1128-
ty_span: this.lower_span(parameter.ty_span),
1129-
span: this.lower_span(parameter.span),
1130-
};
1053+
let (parameters, expr) = this.lower_coroutine_body_with_moved_arguments(
1054+
decl,
1055+
body,
1056+
coroutine_kind,
1057+
CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
1058+
);
11311059

1132-
if is_simple_parameter {
1133-
// If this is the simple case, then we only insert one statement that is
1134-
// `let <pat> = <pat>;`. We re-use the original argument's pattern so that
1135-
// `HirId`s are densely assigned.
1136-
let expr = this.expr_ident(desugared_span, ident, new_parameter_id);
1137-
let stmt = this.stmt_let_pat(
1138-
stmt_attrs,
1139-
desugared_span,
1140-
Some(expr),
1141-
parameter.pat,
1142-
hir::LocalSource::AsyncFn,
1143-
);
1144-
statements.push(stmt);
1145-
} else {
1146-
// If this is not the simple case, then we construct two statements:
1147-
//
1148-
// ```
1149-
// let __argN = __argN;
1150-
// let <pat> = __argN;
1151-
// ```
1152-
//
1153-
// The first statement moves the parameter into the closure and thus ensures
1154-
// that the drop order is correct.
1155-
//
1156-
// The second statement creates the bindings that the user wrote.
1157-
1158-
// Construct the `let mut __argN = __argN;` statement. It must be a mut binding
1159-
// because the user may have specified a `ref mut` binding in the next
1160-
// statement.
1161-
let (move_pat, move_id) = this.pat_ident_binding_mode(
1162-
desugared_span,
1163-
ident,
1164-
hir::BindingAnnotation::MUT,
1165-
);
1166-
let move_expr = this.expr_ident(desugared_span, ident, new_parameter_id);
1167-
let move_stmt = this.stmt_let_pat(
1168-
None,
1169-
desugared_span,
1170-
Some(move_expr),
1171-
move_pat,
1172-
hir::LocalSource::AsyncFn,
1173-
);
1060+
// FIXME(async_fn_track_caller): Can this be moved above?
1061+
let hir_id = this.lower_node_id(coroutine_kind.closure_id());
1062+
this.maybe_forward_track_caller(body.span, fn_id, hir_id);
11741063

1175-
// Construct the `let <pat> = __argN;` statement. We re-use the original
1176-
// parameter's pattern so that `HirId`s are densely assigned.
1177-
let pattern_expr = this.expr_ident(desugared_span, ident, move_id);
1178-
let pattern_stmt = this.stmt_let_pat(
1179-
stmt_attrs,
1180-
desugared_span,
1181-
Some(pattern_expr),
1182-
parameter.pat,
1183-
hir::LocalSource::AsyncFn,
1184-
);
1064+
(parameters, expr)
1065+
})
1066+
}
11851067

1186-
statements.push(move_stmt);
1187-
statements.push(pattern_stmt);
1188-
};
1068+
/// Lowers a desugared coroutine body after moving all of the arguments
1069+
/// into the body. This makes sure that
1070+
fn lower_coroutine_body_with_moved_arguments(
1071+
&mut self,
1072+
decl: &FnDecl,
1073+
body: &Block,
1074+
coroutine_kind: CoroutineKind,
1075+
capture_clause: CaptureBy,
1076+
) -> (&'hir [hir::Param<'hir>], hir::Expr<'hir>) {
1077+
let mut parameters: Vec<hir::Param<'_>> = Vec::new();
1078+
let mut statements: Vec<hir::Stmt<'_>> = Vec::new();
1079+
1080+
// Async function parameters are lowered into the closure body so that they are
1081+
// captured and so that the drop order matches the equivalent non-async functions.
1082+
//
1083+
// from:
1084+
//
1085+
// async fn foo(<pattern>: <ty>, <pattern>: <ty>, <pattern>: <ty>) {
1086+
// <body>
1087+
// }
1088+
//
1089+
// into:
1090+
//
1091+
// fn foo(__arg0: <ty>, __arg1: <ty>, __arg2: <ty>) {
1092+
// async move {
1093+
// let __arg2 = __arg2;
1094+
// let <pattern> = __arg2;
1095+
// let __arg1 = __arg1;
1096+
// let <pattern> = __arg1;
1097+
// let __arg0 = __arg0;
1098+
// let <pattern> = __arg0;
1099+
// drop-temps { <body> } // see comments later in fn for details
1100+
// }
1101+
// }
1102+
//
1103+
// If `<pattern>` is a simple ident, then it is lowered to a single
1104+
// `let <pattern> = <pattern>;` statement as an optimization.
1105+
//
1106+
// Note that the body is embedded in `drop-temps`; an
1107+
// equivalent desugaring would be `return { <body>
1108+
// };`. The key point is that we wish to drop all the
1109+
// let-bound variables and temporaries created in the body
1110+
// (and its tail expression!) before we drop the
1111+
// parameters (c.f. rust-lang/rust#64512).
1112+
for (index, parameter) in decl.inputs.iter().enumerate() {
1113+
let parameter = self.lower_param(parameter);
1114+
let span = parameter.pat.span;
1115+
1116+
// Check if this is a binding pattern, if so, we can optimize and avoid adding a
1117+
// `let <pat> = __argN;` statement. In this case, we do not rename the parameter.
1118+
let (ident, is_simple_parameter) = match parameter.pat.kind {
1119+
hir::PatKind::Binding(hir::BindingAnnotation(ByRef::No, _), _, ident, _) => {
1120+
(ident, true)
1121+
}
1122+
// For `ref mut` or wildcard arguments, we can't reuse the binding, but
1123+
// we can keep the same name for the parameter.
1124+
// This lets rustdoc render it correctly in documentation.
1125+
hir::PatKind::Binding(_, _, ident, _) => (ident, false),
1126+
hir::PatKind::Wild => {
1127+
(Ident::with_dummy_span(rustc_span::symbol::kw::Underscore), false)
1128+
}
1129+
_ => {
1130+
// Replace the ident for bindings that aren't simple.
1131+
let name = format!("__arg{index}");
1132+
let ident = Ident::from_str(&name);
11891133

1190-
parameters.push(new_parameter);
1191-
}
1134+
(ident, false)
1135+
}
1136+
};
11921137

1193-
let mkbody = |this: &mut LoweringContext<'_, 'hir>| {
1194-
// Create a block from the user's function body:
1195-
let user_body = this.lower_block_expr(body);
1138+
let desugared_span = self.mark_span_with_reason(DesugaringKind::Async, span, None);
11961139

1197-
// Transform into `drop-temps { <user-body> }`, an expression:
1198-
let desugared_span =
1199-
this.mark_span_with_reason(DesugaringKind::Async, user_body.span, None);
1200-
let user_body = this.expr_drop_temps(desugared_span, this.arena.alloc(user_body));
1140+
// Construct a parameter representing `__argN: <ty>` to replace the parameter of the
1141+
// async function.
1142+
//
1143+
// If this is the simple case, this parameter will end up being the same as the
1144+
// original parameter, but with a different pattern id.
1145+
let stmt_attrs = self.attrs.get(&parameter.hir_id.local_id).copied();
1146+
let (new_parameter_pat, new_parameter_id) = self.pat_ident(desugared_span, ident);
1147+
let new_parameter = hir::Param {
1148+
hir_id: parameter.hir_id,
1149+
pat: new_parameter_pat,
1150+
ty_span: self.lower_span(parameter.ty_span),
1151+
span: self.lower_span(parameter.span),
1152+
};
12011153

1202-
// As noted above, create the final block like
1154+
if is_simple_parameter {
1155+
// If this is the simple case, then we only insert one statement that is
1156+
// `let <pat> = <pat>;`. We re-use the original argument's pattern so that
1157+
// `HirId`s are densely assigned.
1158+
let expr = self.expr_ident(desugared_span, ident, new_parameter_id);
1159+
let stmt = self.stmt_let_pat(
1160+
stmt_attrs,
1161+
desugared_span,
1162+
Some(expr),
1163+
parameter.pat,
1164+
hir::LocalSource::AsyncFn,
1165+
);
1166+
statements.push(stmt);
1167+
} else {
1168+
// If this is not the simple case, then we construct two statements:
12031169
//
12041170
// ```
1205-
// {
1206-
// let $param_pattern = $raw_param;
1207-
// ...
1208-
// drop-temps { <user-body> }
1209-
// }
1171+
// let __argN = __argN;
1172+
// let <pat> = __argN;
12101173
// ```
1211-
let body = this.block_all(
1174+
//
1175+
// The first statement moves the parameter into the closure and thus ensures
1176+
// that the drop order is correct.
1177+
//
1178+
// The second statement creates the bindings that the user wrote.
1179+
1180+
// Construct the `let mut __argN = __argN;` statement. It must be a mut binding
1181+
// because the user may have specified a `ref mut` binding in the next
1182+
// statement.
1183+
let (move_pat, move_id) =
1184+
self.pat_ident_binding_mode(desugared_span, ident, hir::BindingAnnotation::MUT);
1185+
let move_expr = self.expr_ident(desugared_span, ident, new_parameter_id);
1186+
let move_stmt = self.stmt_let_pat(
1187+
None,
12121188
desugared_span,
1213-
this.arena.alloc_from_iter(statements),
1214-
Some(user_body),
1189+
Some(move_expr),
1190+
move_pat,
1191+
hir::LocalSource::AsyncFn,
12151192
);
12161193

1217-
this.expr_block(body)
1218-
};
1219-
let desugaring_kind = match coroutine_kind {
1220-
CoroutineKind::Async { .. } => hir::CoroutineDesugaring::Async,
1221-
CoroutineKind::Gen { .. } => hir::CoroutineDesugaring::Gen,
1222-
CoroutineKind::AsyncGen { .. } => hir::CoroutineDesugaring::AsyncGen,
1194+
// Construct the `let <pat> = __argN;` statement. We re-use the original
1195+
// parameter's pattern so that `HirId`s are densely assigned.
1196+
let pattern_expr = self.expr_ident(desugared_span, ident, move_id);
1197+
let pattern_stmt = self.stmt_let_pat(
1198+
stmt_attrs,
1199+
desugared_span,
1200+
Some(pattern_expr),
1201+
parameter.pat,
1202+
hir::LocalSource::AsyncFn,
1203+
);
1204+
1205+
statements.push(move_stmt);
1206+
statements.push(pattern_stmt);
12231207
};
1224-
let coroutine_expr = this.make_desugared_coroutine_expr(
1225-
CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
1226-
closure_id,
1227-
None,
1228-
body.span,
1229-
desugaring_kind,
1230-
hir::CoroutineSource::Fn,
1231-
mkbody,
1208+
1209+
parameters.push(new_parameter);
1210+
}
1211+
1212+
let mkbody = |this: &mut LoweringContext<'_, 'hir>| {
1213+
// Create a block from the user's function body:
1214+
let user_body = this.lower_block_expr(body);
1215+
1216+
// Transform into `drop-temps { <user-body> }`, an expression:
1217+
let desugared_span =
1218+
this.mark_span_with_reason(DesugaringKind::Async, user_body.span, None);
1219+
let user_body = this.expr_drop_temps(desugared_span, this.arena.alloc(user_body));
1220+
1221+
// As noted above, create the final block like
1222+
//
1223+
// ```
1224+
// {
1225+
// let $param_pattern = $raw_param;
1226+
// ...
1227+
// drop-temps { <user-body> }
1228+
// }
1229+
// ```
1230+
let body = this.block_all(
1231+
desugared_span,
1232+
this.arena.alloc_from_iter(statements),
1233+
Some(user_body),
12321234
);
12331235

1234-
let hir_id = this.lower_node_id(closure_id);
1235-
this.maybe_forward_track_caller(body.span, fn_id, hir_id);
1236-
let expr = hir::Expr { hir_id, kind: coroutine_expr, span: this.lower_span(body.span) };
1236+
this.expr_block(body)
1237+
};
1238+
let desugaring_kind = match coroutine_kind {
1239+
CoroutineKind::Async { .. } => hir::CoroutineDesugaring::Async,
1240+
CoroutineKind::Gen { .. } => hir::CoroutineDesugaring::Gen,
1241+
CoroutineKind::AsyncGen { .. } => hir::CoroutineDesugaring::AsyncGen,
1242+
};
1243+
let closure_id = coroutine_kind.closure_id();
1244+
let coroutine_expr = self.make_desugared_coroutine_expr(
1245+
capture_clause,
1246+
closure_id,
1247+
None,
1248+
body.span,
1249+
desugaring_kind,
1250+
hir::CoroutineSource::Fn,
1251+
mkbody,
1252+
);
12371253

1238-
(this.arena.alloc_from_iter(parameters), expr)
1239-
})
1254+
let expr = hir::Expr {
1255+
hir_id: self.lower_node_id(closure_id),
1256+
kind: coroutine_expr,
1257+
span: self.lower_span(body.span),
1258+
};
1259+
1260+
(self.arena.alloc_from_iter(parameters), expr)
12401261
}
12411262

12421263
fn lower_method_sig(

0 commit comments

Comments
 (0)