Skip to content

Commit 208dd20

Browse files
borsbors
committed
Auto merge of rust-lang#118847 - eholk:for-await, r=compiler-errors
Add support for `for await` loops This adds support for `for await` loops. This includes parsing, desugaring in AST->HIR lowering, and adding some support functions to the library. Given a loop like: ```rust for await i in iter { ... } ``` this is desugared to something like: ```rust let mut iter = iter.into_async_iter(); while let Some(i) = loop { match core::pin::Pin::new(&mut iter).poll_next(cx) { Poll::Ready(i) => break i, Poll::Pending => yield, } } { ... } ``` This PR also adds a basic `IntoAsyncIterator` trait. This is partly for symmetry with the way `Iterator` and `IntoIterator` work. The other reason is that for async iterators it's helpful to have a place apart from the data structure being iterated over to store state. `IntoAsyncIterator` gives us a good place to do this. I've gated this feature behind `async_for_loop` and opened rust-lang#118898 as the feature tracking issue. r? `@compiler-errors`
2 parents c1fc1d1 + 397f4a1 commit 208dd20

File tree

34 files changed

+367
-79
lines changed

34 files changed

+367
-79
lines changed

Diff for: ‎compiler/rustc_ast/src/ast.rs

+10-3
Original file line numberDiff line numberDiff line change
@@ -1249,7 +1249,7 @@ impl Expr {
12491249
ExprKind::Let(..) => ExprPrecedence::Let,
12501250
ExprKind::If(..) => ExprPrecedence::If,
12511251
ExprKind::While(..) => ExprPrecedence::While,
1252-
ExprKind::ForLoop(..) => ExprPrecedence::ForLoop,
1252+
ExprKind::ForLoop { .. } => ExprPrecedence::ForLoop,
12531253
ExprKind::Loop(..) => ExprPrecedence::Loop,
12541254
ExprKind::Match(..) => ExprPrecedence::Match,
12551255
ExprKind::Closure(..) => ExprPrecedence::Closure,
@@ -1411,10 +1411,10 @@ pub enum ExprKind {
14111411
While(P<Expr>, P<Block>, Option<Label>),
14121412
/// A `for` loop, with an optional label.
14131413
///
1414-
/// `'label: for pat in expr { block }`
1414+
/// `'label: for await? pat in iter { block }`
14151415
///
14161416
/// This is desugared to a combination of `loop` and `match` expressions.
1417-
ForLoop(P<Pat>, P<Expr>, P<Block>, Option<Label>),
1417+
ForLoop { pat: P<Pat>, iter: P<Expr>, body: P<Block>, label: Option<Label>, kind: ForLoopKind },
14181418
/// Conditionless loop (can be exited with `break`, `continue`, or `return`).
14191419
///
14201420
/// `'label: loop { block }`
@@ -1517,6 +1517,13 @@ pub enum ExprKind {
15171517
Err,
15181518
}
15191519

1520+
/// Used to differentiate between `for` loops and `for await` loops.
1521+
#[derive(Clone, Copy, Encodable, Decodable, Debug, PartialEq, Eq)]
1522+
pub enum ForLoopKind {
1523+
For,
1524+
ForAwait,
1525+
}
1526+
15201527
/// Used to differentiate between `async {}` blocks and `gen {}` blocks.
15211528
#[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
15221529
pub enum GenBlockKind {

Diff for: ‎compiler/rustc_ast/src/mut_visit.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -1389,7 +1389,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
13891389
vis.visit_block(body);
13901390
visit_opt(label, |label| vis.visit_label(label));
13911391
}
1392-
ExprKind::ForLoop(pat, iter, body, label) => {
1392+
ExprKind::ForLoop { pat, iter, body, label, kind: _ } => {
13931393
vis.visit_pat(pat);
13941394
vis.visit_expr(iter);
13951395
vis.visit_block(body);

Diff for: ‎compiler/rustc_ast/src/util/classify.rs

+11-3
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ pub fn expr_requires_semi_to_be_stmt(e: &ast::Expr) -> bool {
1919
| ast::ExprKind::Block(..)
2020
| ast::ExprKind::While(..)
2121
| ast::ExprKind::Loop(..)
22-
| ast::ExprKind::ForLoop(..)
22+
| ast::ExprKind::ForLoop { .. }
2323
| ast::ExprKind::TryBlock(..)
2424
| ast::ExprKind::ConstBlock(..)
2525
)
@@ -48,8 +48,16 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<&ast::Expr> {
4848
Closure(closure) => {
4949
expr = &closure.body;
5050
}
51-
Gen(..) | Block(..) | ForLoop(..) | If(..) | Loop(..) | Match(..) | Struct(..)
52-
| TryBlock(..) | While(..) | ConstBlock(_) => break Some(expr),
51+
Gen(..)
52+
| Block(..)
53+
| ForLoop { .. }
54+
| If(..)
55+
| Loop(..)
56+
| Match(..)
57+
| Struct(..)
58+
| TryBlock(..)
59+
| While(..)
60+
| ConstBlock(_) => break Some(expr),
5361

5462
// FIXME: These can end in `}`, but changing these would break stable code.
5563
InlineAsm(_) | OffsetOf(_, _) | MacCall(_) | IncludedBytes(_) | FormatArgs(_) => {

Diff for: ‎compiler/rustc_ast/src/visit.rs

+5-5
Original file line numberDiff line numberDiff line change
@@ -844,11 +844,11 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
844844
visitor.visit_expr(subexpression);
845845
visitor.visit_block(block);
846846
}
847-
ExprKind::ForLoop(pattern, subexpression, block, opt_label) => {
848-
walk_list!(visitor, visit_label, opt_label);
849-
visitor.visit_pat(pattern);
850-
visitor.visit_expr(subexpression);
851-
visitor.visit_block(block);
847+
ExprKind::ForLoop { pat, iter, body, label, kind: _ } => {
848+
walk_list!(visitor, visit_label, label);
849+
visitor.visit_pat(pat);
850+
visitor.visit_expr(iter);
851+
visitor.visit_block(body);
852852
}
853853
ExprKind::Loop(block, opt_label, _) => {
854854
walk_list!(visitor, visit_label, opt_label);

Diff for: ‎compiler/rustc_ast_lowering/src/expr.rs

+95-30
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,12 @@ impl<'hir> LoweringContext<'_, 'hir> {
5656
return ex;
5757
}
5858
// Desugar `ExprForLoop`
59-
// from: `[opt_ident]: for <pat> in <head> <body>`
59+
// from: `[opt_ident]: for await? <pat> in <iter> <body>`
6060
//
6161
// This also needs special handling because the HirId of the returned `hir::Expr` will not
6262
// correspond to the `e.id`, so `lower_expr_for` handles attribute lowering itself.
63-
ExprKind::ForLoop(pat, head, body, opt_label) => {
64-
return self.lower_expr_for(e, pat, head, body, *opt_label);
63+
ExprKind::ForLoop { pat, iter, body, label, kind } => {
64+
return self.lower_expr_for(e, pat, iter, body, *label, *kind);
6565
}
6666
_ => (),
6767
}
@@ -337,7 +337,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
337337
),
338338
ExprKind::Try(sub_expr) => self.lower_expr_try(e.span, sub_expr),
339339

340-
ExprKind::Paren(_) | ExprKind::ForLoop(..) => {
340+
ExprKind::Paren(_) | ExprKind::ForLoop { .. } => {
341341
unreachable!("already handled")
342342
}
343343

@@ -874,6 +874,17 @@ impl<'hir> LoweringContext<'_, 'hir> {
874874
/// }
875875
/// ```
876876
fn lower_expr_await(&mut self, await_kw_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
877+
let expr = self.arena.alloc(self.lower_expr_mut(expr));
878+
self.make_lowered_await(await_kw_span, expr, FutureKind::Future)
879+
}
880+
881+
/// Takes an expr that has already been lowered and generates a desugared await loop around it
882+
fn make_lowered_await(
883+
&mut self,
884+
await_kw_span: Span,
885+
expr: &'hir hir::Expr<'hir>,
886+
await_kind: FutureKind,
887+
) -> hir::ExprKind<'hir> {
877888
let full_span = expr.span.to(await_kw_span);
878889

879890
let is_async_gen = match self.coroutine_kind {
@@ -887,13 +898,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
887898
}
888899
};
889900

890-
let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, None);
901+
let features = match await_kind {
902+
FutureKind::Future => None,
903+
FutureKind::AsyncIterator => Some(self.allow_for_await.clone()),
904+
};
905+
let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, features);
891906
let gen_future_span = self.mark_span_with_reason(
892907
DesugaringKind::Await,
893908
full_span,
894909
Some(self.allow_gen_future.clone()),
895910
);
896-
let expr = self.lower_expr_mut(expr);
897911
let expr_hir_id = expr.hir_id;
898912

899913
// Note that the name of this binding must not be changed to something else because
@@ -934,11 +948,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
934948
hir::LangItem::GetContext,
935949
arena_vec![self; task_context],
936950
);
937-
let call = self.expr_call_lang_item_fn(
938-
span,
939-
hir::LangItem::FuturePoll,
940-
arena_vec![self; new_unchecked, get_context],
941-
);
951+
let call = match await_kind {
952+
FutureKind::Future => self.expr_call_lang_item_fn(
953+
span,
954+
hir::LangItem::FuturePoll,
955+
arena_vec![self; new_unchecked, get_context],
956+
),
957+
FutureKind::AsyncIterator => self.expr_call_lang_item_fn(
958+
span,
959+
hir::LangItem::AsyncIteratorPollNext,
960+
arena_vec![self; new_unchecked, get_context],
961+
),
962+
};
942963
self.arena.alloc(self.expr_unsafe(call))
943964
};
944965

@@ -1020,11 +1041,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
10201041
let awaitee_arm = self.arm(awaitee_pat, loop_expr);
10211042

10221043
// `match ::std::future::IntoFuture::into_future(<expr>) { ... }`
1023-
let into_future_expr = self.expr_call_lang_item_fn(
1024-
span,
1025-
hir::LangItem::IntoFutureIntoFuture,
1026-
arena_vec![self; expr],
1027-
);
1044+
let into_future_expr = match await_kind {
1045+
FutureKind::Future => self.expr_call_lang_item_fn(
1046+
span,
1047+
hir::LangItem::IntoFutureIntoFuture,
1048+
arena_vec![self; *expr],
1049+
),
1050+
// Not needed for `for await` because we expect to have already called
1051+
// `IntoAsyncIterator::into_async_iter` on it.
1052+
FutureKind::AsyncIterator => expr,
1053+
};
10281054

10291055
// match <into_future_expr> {
10301056
// mut __awaitee => loop { .. }
@@ -1685,6 +1711,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
16851711
head: &Expr,
16861712
body: &Block,
16871713
opt_label: Option<Label>,
1714+
loop_kind: ForLoopKind,
16881715
) -> hir::Expr<'hir> {
16891716
let head = self.lower_expr_mut(head);
16901717
let pat = self.lower_pat(pat);
@@ -1713,17 +1740,41 @@ impl<'hir> LoweringContext<'_, 'hir> {
17131740
let (iter_pat, iter_pat_nid) =
17141741
self.pat_ident_binding_mode(head_span, iter, hir::BindingAnnotation::MUT);
17151742

1716-
// `match Iterator::next(&mut iter) { ... }`
17171743
let match_expr = {
17181744
let iter = self.expr_ident(head_span, iter, iter_pat_nid);
1719-
let ref_mut_iter = self.expr_mut_addr_of(head_span, iter);
1720-
let next_expr = self.expr_call_lang_item_fn(
1721-
head_span,
1722-
hir::LangItem::IteratorNext,
1723-
arena_vec![self; ref_mut_iter],
1724-
);
1745+
let next_expr = match loop_kind {
1746+
ForLoopKind::For => {
1747+
// `Iterator::next(&mut iter)`
1748+
let ref_mut_iter = self.expr_mut_addr_of(head_span, iter);
1749+
self.expr_call_lang_item_fn(
1750+
head_span,
1751+
hir::LangItem::IteratorNext,
1752+
arena_vec![self; ref_mut_iter],
1753+
)
1754+
}
1755+
ForLoopKind::ForAwait => {
1756+
// we'll generate `unsafe { Pin::new_unchecked(&mut iter) })` and then pass this
1757+
// to make_lowered_await with `FutureKind::AsyncIterator` which will generator
1758+
// calls to `poll_next`. In user code, this would probably be a call to
1759+
// `Pin::as_mut` but here it's easy enough to do `new_unchecked`.
1760+
1761+
// `&mut iter`
1762+
let iter = self.expr_mut_addr_of(head_span, iter);
1763+
// `Pin::new_unchecked(...)`
1764+
let iter = self.arena.alloc(self.expr_call_lang_item_fn_mut(
1765+
head_span,
1766+
hir::LangItem::PinNewUnchecked,
1767+
arena_vec![self; iter],
1768+
));
1769+
// `unsafe { ... }`
1770+
let iter = self.arena.alloc(self.expr_unsafe(iter));
1771+
let kind = self.make_lowered_await(head_span, iter, FutureKind::AsyncIterator);
1772+
self.arena.alloc(hir::Expr { hir_id: self.next_id(), kind, span: head_span })
1773+
}
1774+
};
17251775
let arms = arena_vec![self; none_arm, some_arm];
17261776

1777+
// `match $next_expr { ... }`
17271778
self.expr_match(head_span, next_expr, arms, hir::MatchSource::ForLoopDesugar)
17281779
};
17291780
let match_stmt = self.stmt_expr(for_span, match_expr);
@@ -1743,13 +1794,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
17431794
// `mut iter => { ... }`
17441795
let iter_arm = self.arm(iter_pat, loop_expr);
17451796

1746-
// `match ::std::iter::IntoIterator::into_iter(<head>) { ... }`
1747-
let into_iter_expr = {
1748-
self.expr_call_lang_item_fn(
1749-
head_span,
1750-
hir::LangItem::IntoIterIntoIter,
1751-
arena_vec![self; head],
1752-
)
1797+
let into_iter_expr = match loop_kind {
1798+
ForLoopKind::For => {
1799+
// `::std::iter::IntoIterator::into_iter(<head>)`
1800+
self.expr_call_lang_item_fn(
1801+
head_span,
1802+
hir::LangItem::IntoIterIntoIter,
1803+
arena_vec![self; head],
1804+
)
1805+
}
1806+
ForLoopKind::ForAwait => self.arena.alloc(head),
17531807
};
17541808

17551809
let match_expr = self.arena.alloc(self.expr_match(
@@ -2152,3 +2206,14 @@ impl<'hir> LoweringContext<'_, 'hir> {
21522206
}
21532207
}
21542208
}
2209+
2210+
/// Used by [`LoweringContext::make_lowered_await`] to customize the desugaring based on what kind
2211+
/// of future we are awaiting.
2212+
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
2213+
enum FutureKind {
2214+
/// We are awaiting a normal future
2215+
Future,
2216+
/// We are awaiting something that's known to be an AsyncIterator (i.e. we are in the header of
2217+
/// a `for await` loop)
2218+
AsyncIterator,
2219+
}

Diff for: ‎compiler/rustc_ast_lowering/src/lib.rs

+2
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ struct LoweringContext<'a, 'hir> {
130130
allow_try_trait: Lrc<[Symbol]>,
131131
allow_gen_future: Lrc<[Symbol]>,
132132
allow_async_iterator: Lrc<[Symbol]>,
133+
allow_for_await: Lrc<[Symbol]>,
133134

134135
/// Mapping from generics `def_id`s to TAIT generics `def_id`s.
135136
/// For each captured lifetime (e.g., 'a), we create a new lifetime parameter that is a generic
@@ -174,6 +175,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
174175
} else {
175176
[sym::gen_future].into()
176177
},
178+
allow_for_await: [sym::async_iterator].into(),
177179
// FIXME(gen_blocks): how does `closure_track_caller`/`async_fn_track_caller`
178180
// interact with `gen`/`async gen` blocks
179181
allow_async_iterator: [sym::gen_future, sym::async_iterator].into(),

Diff for: ‎compiler/rustc_ast_passes/src/feature_gate.rs

+1
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,7 @@ pub fn check_crate(krate: &ast::Crate, sess: &Session, features: &Features) {
527527
"async closures are unstable",
528528
"to use an async block, remove the `||`: `async {`"
529529
);
530+
gate_all!(async_for_loop, "`for await` loops are experimental");
530531
gate_all!(
531532
closure_lifetime_binder,
532533
"`for<...>` binders for closures are experimental",

Diff for: ‎compiler/rustc_ast_pretty/src/pprust/state/expr.rs

+7-3
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use crate::pp::Breaks::Inconsistent;
22
use crate::pprust::state::{AnnNode, PrintState, State, INDENT_UNIT};
3+
use ast::ForLoopKind;
34
use itertools::{Itertools, Position};
45
use rustc_ast::ptr::P;
56
use rustc_ast::token;
@@ -418,20 +419,23 @@ impl<'a> State<'a> {
418419
self.space();
419420
self.print_block_with_attrs(blk, attrs);
420421
}
421-
ast::ExprKind::ForLoop(pat, iter, blk, opt_label) => {
422-
if let Some(label) = opt_label {
422+
ast::ExprKind::ForLoop { pat, iter, body, label, kind } => {
423+
if let Some(label) = label {
423424
self.print_ident(label.ident);
424425
self.word_space(":");
425426
}
426427
self.cbox(0);
427428
self.ibox(0);
428429
self.word_nbsp("for");
430+
if kind == &ForLoopKind::ForAwait {
431+
self.word_nbsp("await");
432+
}
429433
self.print_pat(pat);
430434
self.space();
431435
self.word_space("in");
432436
self.print_expr_as_cond(iter);
433437
self.space();
434-
self.print_block_with_attrs(blk, attrs);
438+
self.print_block_with_attrs(body, attrs);
435439
}
436440
ast::ExprKind::Loop(blk, opt_label, _) => {
437441
if let Some(label) = opt_label {

Diff for: ‎compiler/rustc_builtin_macros/src/assert/context.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
303303
| ExprKind::Continue(_)
304304
| ExprKind::Err
305305
| ExprKind::Field(_, _)
306-
| ExprKind::ForLoop(_, _, _, _)
306+
| ExprKind::ForLoop { .. }
307307
| ExprKind::FormatArgs(_)
308308
| ExprKind::IncludedBytes(..)
309309
| ExprKind::InlineAsm(_)

Diff for: ‎compiler/rustc_feature/src/unstable.rs

+2
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,8 @@ declare_features! (
357357
(unstable, async_closure, "1.37.0", Some(62290)),
358358
/// Allows `#[track_caller]` on async functions.
359359
(unstable, async_fn_track_caller, "1.73.0", Some(110011)),
360+
/// Allows `for await` loops.
361+
(unstable, async_for_loop, "CURRENT_RUSTC_VERSION", Some(118898)),
360362
/// Allows builtin # foo() syntax
361363
(unstable, builtin_syntax, "1.71.0", Some(110680)),
362364
/// Treat `extern "C"` function as nounwind.

Diff for: ‎compiler/rustc_hir/src/lang_items.rs

+2
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,8 @@ language_item_table! {
307307
Context, sym::Context, context, Target::Struct, GenericRequirement::None;
308308
FuturePoll, sym::poll, future_poll_fn, Target::Method(MethodKind::Trait { body: false }), GenericRequirement::None;
309309

310+
AsyncIteratorPollNext, sym::async_iterator_poll_next, async_iterator_poll_next, Target::Method(MethodKind::Trait { body: false }), GenericRequirement::Exact(0);
311+
310312
Option, sym::Option, option_type, Target::Enum, GenericRequirement::None;
311313
OptionSome, sym::Some, option_some_variant, Target::Variant, GenericRequirement::None;
312314
OptionNone, sym::None, option_none_variant, Target::Variant, GenericRequirement::None;

Diff for: ‎compiler/rustc_lint/src/unused.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -852,8 +852,8 @@ trait UnusedDelimLint {
852852
(cond, UnusedDelimsCtx::WhileCond, true, Some(left), Some(right), true)
853853
}
854854

855-
ForLoop(_, ref cond, ref block, ..) => {
856-
(cond, UnusedDelimsCtx::ForIterExpr, true, None, Some(block.span.lo()), true)
855+
ForLoop { ref iter, ref body, .. } => {
856+
(iter, UnusedDelimsCtx::ForIterExpr, true, None, Some(body.span.lo()), true)
857857
}
858858

859859
Match(ref head, _) if Self::LINT_EXPR_IN_PATTERN_MATCHING_CTX => {
@@ -1085,7 +1085,7 @@ impl EarlyLintPass for UnusedParens {
10851085
}
10861086

10871087
match e.kind {
1088-
ExprKind::Let(ref pat, _, _, _) | ExprKind::ForLoop(ref pat, ..) => {
1088+
ExprKind::Let(ref pat, _, _, _) | ExprKind::ForLoop { ref pat, .. } => {
10891089
self.check_unused_parens_pat(cx, pat, false, false, (true, true));
10901090
}
10911091
// We ignore parens in cases like `if (((let Some(0) = Some(1))))` because we already

0 commit comments

Comments
 (0)