Skip to content

Commit 8bb1eae

Browse files
compiler-errorscompiler-errors
committed
Introduce helper that deals with moving async args into the coroutine
1 parent fa0dc20 commit 8bb1eae

File tree

1 file changed

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

1 file changed

+193
-170
lines changed

‎compiler/rustc_ast_lowering/src/item.rs

+193-170
Original file line numberDiff line numberDiff line change
@@ -1082,194 +1082,217 @@ impl<'hir> LoweringContext<'_, 'hir> {
10821082
let (Some(coroutine_kind), Some(body)) = (coroutine_kind, body) else {
10831083
return self.lower_fn_body_block(span, decl, body);
10841084
};
1085-
let closure_id = coroutine_kind.closure_id();
1086-
10871085
self.lower_body(|this| {
1088-
let mut parameters: Vec<hir::Param<'_>> = Vec::new();
1089-
let mut statements: Vec<hir::Stmt<'_>> = Vec::new();
1090-
1091-
// Async function parameters are lowered into the closure body so that they are
1092-
// captured and so that the drop order matches the equivalent non-async functions.
1093-
//
1094-
// from:
1095-
//
1096-
// async fn foo(<pattern>: <ty>, <pattern>: <ty>, <pattern>: <ty>) {
1097-
// <body>
1098-
// }
1099-
//
1100-
// into:
1101-
//
1102-
// fn foo(__arg0: <ty>, __arg1: <ty>, __arg2: <ty>) {
1103-
// async move {
1104-
// let __arg2 = __arg2;
1105-
// let <pattern> = __arg2;
1106-
// let __arg1 = __arg1;
1107-
// let <pattern> = __arg1;
1108-
// let __arg0 = __arg0;
1109-
// let <pattern> = __arg0;
1110-
// drop-temps { <body> } // see comments later in fn for details
1111-
// }
1112-
// }
1113-
//
1114-
// If `<pattern>` is a simple ident, then it is lowered to a single
1115-
// `let <pattern> = <pattern>;` statement as an optimization.
1116-
//
1117-
// Note that the body is embedded in `drop-temps`; an
1118-
// equivalent desugaring would be `return { <body>
1119-
// };`. The key point is that we wish to drop all the
1120-
// let-bound variables and temporaries created in the body
1121-
// (and its tail expression!) before we drop the
1122-
// parameters (c.f. rust-lang/rust#64512).
1123-
for (index, parameter) in decl.inputs.iter().enumerate() {
1124-
let parameter = this.lower_param(parameter);
1125-
let span = parameter.pat.span;
1126-
1127-
// Check if this is a binding pattern, if so, we can optimize and avoid adding a
1128-
// `let <pat> = __argN;` statement. In this case, we do not rename the parameter.
1129-
let (ident, is_simple_parameter) = match parameter.pat.kind {
1130-
hir::PatKind::Binding(hir::BindingAnnotation(ByRef::No, _), _, ident, _) => {
1131-
(ident, true)
1132-
}
1133-
// For `ref mut` or wildcard arguments, we can't reuse the binding, but
1134-
// we can keep the same name for the parameter.
1135-
// This lets rustdoc render it correctly in documentation.
1136-
hir::PatKind::Binding(_, _, ident, _) => (ident, false),
1137-
hir::PatKind::Wild => {
1138-
(Ident::with_dummy_span(rustc_span::symbol::kw::Underscore), false)
1139-
}
1140-
_ => {
1141-
// Replace the ident for bindings that aren't simple.
1142-
let name = format!("__arg{index}");
1143-
let ident = Ident::from_str(&name);
1144-
1145-
(ident, false)
1146-
}
1147-
};
1148-
1149-
let desugared_span = this.mark_span_with_reason(DesugaringKind::Async, span, None);
1150-
1151-
// Construct a parameter representing `__argN: <ty>` to replace the parameter of the
1152-
// async function.
1153-
//
1154-
// If this is the simple case, this parameter will end up being the same as the
1155-
// original parameter, but with a different pattern id.
1156-
let stmt_attrs = this.attrs.get(&parameter.hir_id.local_id).copied();
1157-
let (new_parameter_pat, new_parameter_id) = this.pat_ident(desugared_span, ident);
1158-
let new_parameter = hir::Param {
1159-
hir_id: parameter.hir_id,
1160-
pat: new_parameter_pat,
1161-
ty_span: this.lower_span(parameter.ty_span),
1162-
span: this.lower_span(parameter.span),
1163-
};
1086+
let (parameters, expr) = this.lower_coroutine_body_with_moved_arguments(
1087+
decl,
1088+
body,
1089+
coroutine_kind,
1090+
CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
1091+
);
11641092

1165-
if is_simple_parameter {
1166-
// If this is the simple case, then we only insert one statement that is
1167-
// `let <pat> = <pat>;`. We re-use the original argument's pattern so that
1168-
// `HirId`s are densely assigned.
1169-
let expr = this.expr_ident(desugared_span, ident, new_parameter_id);
1170-
let stmt = this.stmt_let_pat(
1171-
stmt_attrs,
1172-
desugared_span,
1173-
Some(expr),
1174-
parameter.pat,
1175-
hir::LocalSource::AsyncFn,
1176-
);
1177-
statements.push(stmt);
1178-
} else {
1179-
// If this is not the simple case, then we construct two statements:
1180-
//
1181-
// ```
1182-
// let __argN = __argN;
1183-
// let <pat> = __argN;
1184-
// ```
1185-
//
1186-
// The first statement moves the parameter into the closure and thus ensures
1187-
// that the drop order is correct.
1188-
//
1189-
// The second statement creates the bindings that the user wrote.
1190-
1191-
// Construct the `let mut __argN = __argN;` statement. It must be a mut binding
1192-
// because the user may have specified a `ref mut` binding in the next
1193-
// statement.
1194-
let (move_pat, move_id) = this.pat_ident_binding_mode(
1195-
desugared_span,
1196-
ident,
1197-
hir::BindingAnnotation::MUT,
1198-
);
1199-
let move_expr = this.expr_ident(desugared_span, ident, new_parameter_id);
1200-
let move_stmt = this.stmt_let_pat(
1201-
None,
1202-
desugared_span,
1203-
Some(move_expr),
1204-
move_pat,
1205-
hir::LocalSource::AsyncFn,
1206-
);
1093+
// FIXME(async_fn_track_caller): Can this be moved above?
1094+
let hir_id = this.lower_node_id(coroutine_kind.closure_id());
1095+
this.maybe_forward_track_caller(body.span, fn_id, hir_id);
12071096

1208-
// Construct the `let <pat> = __argN;` statement. We re-use the original
1209-
// parameter's pattern so that `HirId`s are densely assigned.
1210-
let pattern_expr = this.expr_ident(desugared_span, ident, move_id);
1211-
let pattern_stmt = this.stmt_let_pat(
1212-
stmt_attrs,
1213-
desugared_span,
1214-
Some(pattern_expr),
1215-
parameter.pat,
1216-
hir::LocalSource::AsyncFn,
1217-
);
1097+
(parameters, expr)
1098+
})
1099+
}
12181100

1219-
statements.push(move_stmt);
1220-
statements.push(pattern_stmt);
1221-
};
1101+
/// Lowers a desugared coroutine body after moving all of the arguments
1102+
/// into the body. This is to make sure that the future actually owns the
1103+
/// arguments that are passed to the function, and to ensure things like
1104+
/// drop order are stable.
1105+
fn lower_coroutine_body_with_moved_arguments(
1106+
&mut self,
1107+
decl: &FnDecl,
1108+
body: &Block,
1109+
coroutine_kind: CoroutineKind,
1110+
capture_clause: CaptureBy,
1111+
) -> (&'hir [hir::Param<'hir>], hir::Expr<'hir>) {
1112+
let mut parameters: Vec<hir::Param<'_>> = Vec::new();
1113+
let mut statements: Vec<hir::Stmt<'_>> = Vec::new();
1114+
1115+
// Async function parameters are lowered into the closure body so that they are
1116+
// captured and so that the drop order matches the equivalent non-async functions.
1117+
//
1118+
// from:
1119+
//
1120+
// async fn foo(<pattern>: <ty>, <pattern>: <ty>, <pattern>: <ty>) {
1121+
// <body>
1122+
// }
1123+
//
1124+
// into:
1125+
//
1126+
// fn foo(__arg0: <ty>, __arg1: <ty>, __arg2: <ty>) {
1127+
// async move {
1128+
// let __arg2 = __arg2;
1129+
// let <pattern> = __arg2;
1130+
// let __arg1 = __arg1;
1131+
// let <pattern> = __arg1;
1132+
// let __arg0 = __arg0;
1133+
// let <pattern> = __arg0;
1134+
// drop-temps { <body> } // see comments later in fn for details
1135+
// }
1136+
// }
1137+
//
1138+
// If `<pattern>` is a simple ident, then it is lowered to a single
1139+
// `let <pattern> = <pattern>;` statement as an optimization.
1140+
//
1141+
// Note that the body is embedded in `drop-temps`; an
1142+
// equivalent desugaring would be `return { <body>
1143+
// };`. The key point is that we wish to drop all the
1144+
// let-bound variables and temporaries created in the body
1145+
// (and its tail expression!) before we drop the
1146+
// parameters (c.f. rust-lang/rust#64512).
1147+
for (index, parameter) in decl.inputs.iter().enumerate() {
1148+
let parameter = self.lower_param(parameter);
1149+
let span = parameter.pat.span;
1150+
1151+
// Check if this is a binding pattern, if so, we can optimize and avoid adding a
1152+
// `let <pat> = __argN;` statement. In this case, we do not rename the parameter.
1153+
let (ident, is_simple_parameter) = match parameter.pat.kind {
1154+
hir::PatKind::Binding(hir::BindingAnnotation(ByRef::No, _), _, ident, _) => {
1155+
(ident, true)
1156+
}
1157+
// For `ref mut` or wildcard arguments, we can't reuse the binding, but
1158+
// we can keep the same name for the parameter.
1159+
// This lets rustdoc render it correctly in documentation.
1160+
hir::PatKind::Binding(_, _, ident, _) => (ident, false),
1161+
hir::PatKind::Wild => {
1162+
(Ident::with_dummy_span(rustc_span::symbol::kw::Underscore), false)
1163+
}
1164+
_ => {
1165+
// Replace the ident for bindings that aren't simple.
1166+
let name = format!("__arg{index}");
1167+
let ident = Ident::from_str(&name);
12221168

1223-
parameters.push(new_parameter);
1224-
}
1169+
(ident, false)
1170+
}
1171+
};
12251172

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

1230-
// Transform into `drop-temps { <user-body> }`, an expression:
1231-
let desugared_span =
1232-
this.mark_span_with_reason(DesugaringKind::Async, user_body.span, None);
1233-
let user_body = this.expr_drop_temps(desugared_span, this.arena.alloc(user_body));
1175+
// Construct a parameter representing `__argN: <ty>` to replace the parameter of the
1176+
// async function.
1177+
//
1178+
// If this is the simple case, this parameter will end up being the same as the
1179+
// original parameter, but with a different pattern id.
1180+
let stmt_attrs = self.attrs.get(&parameter.hir_id.local_id).copied();
1181+
let (new_parameter_pat, new_parameter_id) = self.pat_ident(desugared_span, ident);
1182+
let new_parameter = hir::Param {
1183+
hir_id: parameter.hir_id,
1184+
pat: new_parameter_pat,
1185+
ty_span: self.lower_span(parameter.ty_span),
1186+
span: self.lower_span(parameter.span),
1187+
};
12341188

1235-
// As noted above, create the final block like
1189+
if is_simple_parameter {
1190+
// If this is the simple case, then we only insert one statement that is
1191+
// `let <pat> = <pat>;`. We re-use the original argument's pattern so that
1192+
// `HirId`s are densely assigned.
1193+
let expr = self.expr_ident(desugared_span, ident, new_parameter_id);
1194+
let stmt = self.stmt_let_pat(
1195+
stmt_attrs,
1196+
desugared_span,
1197+
Some(expr),
1198+
parameter.pat,
1199+
hir::LocalSource::AsyncFn,
1200+
);
1201+
statements.push(stmt);
1202+
} else {
1203+
// If this is not the simple case, then we construct two statements:
12361204
//
12371205
// ```
1238-
// {
1239-
// let $param_pattern = $raw_param;
1240-
// ...
1241-
// drop-temps { <user-body> }
1242-
// }
1206+
// let __argN = __argN;
1207+
// let <pat> = __argN;
12431208
// ```
1244-
let body = this.block_all(
1209+
//
1210+
// The first statement moves the parameter into the closure and thus ensures
1211+
// that the drop order is correct.
1212+
//
1213+
// The second statement creates the bindings that the user wrote.
1214+
1215+
// Construct the `let mut __argN = __argN;` statement. It must be a mut binding
1216+
// because the user may have specified a `ref mut` binding in the next
1217+
// statement.
1218+
let (move_pat, move_id) =
1219+
self.pat_ident_binding_mode(desugared_span, ident, hir::BindingAnnotation::MUT);
1220+
let move_expr = self.expr_ident(desugared_span, ident, new_parameter_id);
1221+
let move_stmt = self.stmt_let_pat(
1222+
None,
12451223
desugared_span,
1246-
this.arena.alloc_from_iter(statements),
1247-
Some(user_body),
1224+
Some(move_expr),
1225+
move_pat,
1226+
hir::LocalSource::AsyncFn,
12481227
);
12491228

1250-
this.expr_block(body)
1251-
};
1252-
let desugaring_kind = match coroutine_kind {
1253-
CoroutineKind::Async { .. } => hir::CoroutineDesugaring::Async,
1254-
CoroutineKind::Gen { .. } => hir::CoroutineDesugaring::Gen,
1255-
CoroutineKind::AsyncGen { .. } => hir::CoroutineDesugaring::AsyncGen,
1229+
// Construct the `let <pat> = __argN;` statement. We re-use the original
1230+
// parameter's pattern so that `HirId`s are densely assigned.
1231+
let pattern_expr = self.expr_ident(desugared_span, ident, move_id);
1232+
let pattern_stmt = self.stmt_let_pat(
1233+
stmt_attrs,
1234+
desugared_span,
1235+
Some(pattern_expr),
1236+
parameter.pat,
1237+
hir::LocalSource::AsyncFn,
1238+
);
1239+
1240+
statements.push(move_stmt);
1241+
statements.push(pattern_stmt);
12561242
};
1257-
let coroutine_expr = this.make_desugared_coroutine_expr(
1258-
CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
1259-
closure_id,
1260-
None,
1261-
body.span,
1262-
desugaring_kind,
1263-
hir::CoroutineSource::Fn,
1264-
mkbody,
1243+
1244+
parameters.push(new_parameter);
1245+
}
1246+
1247+
let mkbody = |this: &mut LoweringContext<'_, 'hir>| {
1248+
// Create a block from the user's function body:
1249+
let user_body = this.lower_block_expr(body);
1250+
1251+
// Transform into `drop-temps { <user-body> }`, an expression:
1252+
let desugared_span =
1253+
this.mark_span_with_reason(DesugaringKind::Async, user_body.span, None);
1254+
let user_body = this.expr_drop_temps(desugared_span, this.arena.alloc(user_body));
1255+
1256+
// As noted above, create the final block like
1257+
//
1258+
// ```
1259+
// {
1260+
// let $param_pattern = $raw_param;
1261+
// ...
1262+
// drop-temps { <user-body> }
1263+
// }
1264+
// ```
1265+
let body = this.block_all(
1266+
desugared_span,
1267+
this.arena.alloc_from_iter(statements),
1268+
Some(user_body),
12651269
);
12661270

1267-
let hir_id = this.lower_node_id(closure_id);
1268-
this.maybe_forward_track_caller(body.span, fn_id, hir_id);
1269-
let expr = hir::Expr { hir_id, kind: coroutine_expr, span: this.lower_span(body.span) };
1271+
this.expr_block(body)
1272+
};
1273+
let desugaring_kind = match coroutine_kind {
1274+
CoroutineKind::Async { .. } => hir::CoroutineDesugaring::Async,
1275+
CoroutineKind::Gen { .. } => hir::CoroutineDesugaring::Gen,
1276+
CoroutineKind::AsyncGen { .. } => hir::CoroutineDesugaring::AsyncGen,
1277+
};
1278+
let closure_id = coroutine_kind.closure_id();
1279+
let coroutine_expr = self.make_desugared_coroutine_expr(
1280+
capture_clause,
1281+
closure_id,
1282+
None,
1283+
body.span,
1284+
desugaring_kind,
1285+
hir::CoroutineSource::Fn,
1286+
mkbody,
1287+
);
12701288

1271-
(this.arena.alloc_from_iter(parameters), expr)
1272-
})
1289+
let expr = hir::Expr {
1290+
hir_id: self.lower_node_id(closure_id),
1291+
kind: coroutine_expr,
1292+
span: self.lower_span(body.span),
1293+
};
1294+
1295+
(self.arena.alloc_from_iter(parameters), expr)
12731296
}
12741297

12751298
fn lower_method_sig(

0 commit comments

Comments
 (0)