Skip to content

Commit 6d42dbf

Browse files
eholkeholk
committed
Desugar for await loops
1 parent 91622cf commit 6d42dbf

File tree

7 files changed

+87
-15
lines changed

7 files changed

+87
-15
lines changed

‎compiler/rustc_ast/src/mut_visit.rs‎

Lines changed: 1 addition & 1 deletion
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, is_await: _} => {
1392+
ExprKind::ForLoop { pat, iter, body, label, is_await: _ } => {
13931393
vis.visit_pat(pat);
13941394
vis.visit_expr(iter);
13951395
vis.visit_block(body);

‎compiler/rustc_ast_lowering/src/expr.rs‎

Lines changed: 48 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -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

@@ -871,6 +871,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
871871
/// }
872872
/// ```
873873
fn lower_expr_await(&mut self, await_kw_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
874+
let expr = self.lower_expr_mut(expr);
875+
self.make_lowered_await(await_kw_span, &expr)
876+
}
877+
878+
/// Takes an expr that has already been lowered and generates a desugared await loop around it
879+
fn make_lowered_await(
880+
&mut self,
881+
await_kw_span: Span,
882+
expr: &hir::Expr<'hir>,
883+
) -> hir::ExprKind<'hir> {
874884
let full_span = expr.span.to(await_kw_span);
875885

876886
let is_async_gen = match self.coroutine_kind {
@@ -890,7 +900,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
890900
full_span,
891901
Some(self.allow_gen_future.clone()),
892902
);
893-
let expr = self.lower_expr_mut(expr);
894903
let expr_hir_id = expr.hir_id;
895904

896905
// Note that the name of this binding must not be changed to something else because
@@ -1021,7 +1030,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
10211030
let into_future_expr = self.expr_call_lang_item_fn(
10221031
span,
10231032
hir::LangItem::IntoFutureIntoFuture,
1024-
arena_vec![self; expr],
1033+
arena_vec![self; *expr],
10251034
);
10261035

10271036
// match <into_future_expr> {
@@ -1670,7 +1679,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
16701679
head: &Expr,
16711680
body: &Block,
16721681
opt_label: Option<Label>,
1673-
_is_await: bool,
1682+
is_await: bool,
16741683
) -> hir::Expr<'hir> {
16751684
let head = self.lower_expr_mut(head);
16761685
let pat = self.lower_pat(pat);
@@ -1699,15 +1708,33 @@ impl<'hir> LoweringContext<'_, 'hir> {
16991708
let (iter_pat, iter_pat_nid) =
17001709
self.pat_ident_binding_mode(head_span, iter, hir::BindingAnnotation::MUT);
17011710

1702-
// `match Iterator::next(&mut iter) { ... }`
17031711
let match_expr = {
17041712
let iter = self.expr_ident(head_span, iter, iter_pat_nid);
1705-
let ref_mut_iter = self.expr_mut_addr_of(head_span, iter);
1706-
let next_expr = self.expr_call_lang_item_fn(
1707-
head_span,
1708-
hir::LangItem::IteratorNext,
1709-
arena_vec![self; ref_mut_iter],
1710-
);
1713+
let next_expr = if is_await {
1714+
// `async_iter_next(unsafe { Pin::new_unchecked(&mut iter) }).await`
1715+
let future = self.expr_mut_addr_of(head_span, iter);
1716+
let future = self.arena.alloc(self.expr_call_lang_item_fn_mut(
1717+
head_span,
1718+
hir::LangItem::PinNewUnchecked,
1719+
arena_vec![self; future],
1720+
));
1721+
let future = self.expr_unsafe(future);
1722+
let future = self.expr_call_lang_item_fn(
1723+
head_span,
1724+
hir::LangItem::AsyncIterNext,
1725+
arena_vec![self; future],
1726+
);
1727+
let kind = self.make_lowered_await(head_span, future);
1728+
self.arena.alloc(hir::Expr { hir_id: self.next_id(), kind, span: head_span })
1729+
} else {
1730+
// `match Iterator::next(&mut iter) { ... }`
1731+
let ref_mut_iter = self.expr_mut_addr_of(head_span, iter);
1732+
self.expr_call_lang_item_fn(
1733+
head_span,
1734+
hir::LangItem::IteratorNext,
1735+
arena_vec![self; ref_mut_iter],
1736+
)
1737+
};
17111738
let arms = arena_vec![self; none_arm, some_arm];
17121739

17131740
self.expr_match(head_span, next_expr, arms, hir::MatchSource::ForLoopDesugar)
@@ -1729,8 +1756,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
17291756
// `mut iter => { ... }`
17301757
let iter_arm = self.arm(iter_pat, loop_expr);
17311758

1732-
// `match ::std::iter::IntoIterator::into_iter(<head>) { ... }`
1733-
let into_iter_expr = {
1759+
let into_iter_expr = if is_await {
1760+
// `::core::async_iter::IntoAsyncIterator::into_async_iter(<head>)`
1761+
let iter = self.expr_call_lang_item_fn(
1762+
head_span,
1763+
hir::LangItem::IntoAsyncIterIntoIter,
1764+
arena_vec![self; head],
1765+
);
1766+
self.arena.alloc(self.expr_mut_addr_of(head_span, iter))
1767+
} else {
1768+
// `::std::iter::IntoIterator::into_iter(<head>)`
17341769
self.expr_call_lang_item_fn(
17351770
head_span,
17361771
hir::LangItem::IntoIterIntoIter,

‎compiler/rustc_builtin_macros/src/assert/context.rs‎

Lines changed: 1 addition & 1 deletion
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(_)

‎compiler/rustc_hir/src/lang_items.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,7 @@ language_item_table! {
308308
FuturePoll, sym::poll, future_poll_fn, Target::Method(MethodKind::Trait { body: false }), GenericRequirement::None;
309309

310310
AsyncIterNext, sym::async_iter_next, async_iter_next, Target::Fn, GenericRequirement::Exact(2);
311+
IntoAsyncIterIntoIter, sym::into_async_iter_into_iter, into_async_iter_into_iter, Target::Method(MethodKind::Trait { body: false }), GenericRequirement::Exact(1);
311312

312313
Option, sym::Option, option_type, Target::Enum, GenericRequirement::None;
313314
OptionSome, sym::Some, option_some_variant, Target::Variant, GenericRequirement::None;

‎compiler/rustc_span/src/symbol.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -894,6 +894,7 @@ symbols! {
894894
instruction_set,
895895
integer_: "integer", // underscore to avoid clashing with the function `sym::integer` below
896896
integral,
897+
into_async_iter_into_iter,
897898
into_future,
898899
into_iter,
899900
intra_doc_pointers,

‎library/core/src/async_iter/async_iter.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ pub trait IntoAsyncIterator {
144144
type IntoAsyncIter: AsyncIterator<Item = Self::Item>;
145145

146146
/// Converts `self` into an async iterator
147+
#[cfg_attr(not(bootstrap), lang = "into_async_iter_into_iter")]
147148
fn into_async_iter(self) -> Self::IntoAsyncIter;
148149
}
149150

‎tests/ui/async-await/for-await.rs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
// run-pass
2+
// edition: 2021
3+
#![feature(async_iterator, async_iter_from_iter, const_waker)]
4+
5+
use std::future::Future;
6+
7+
// make sure a simple for await loop works
8+
async fn real_main() {
9+
let iter = core::async_iter::from_iter(0..3);
10+
let mut count = 0;
11+
for await i in iter {
12+
assert_eq!(i, count);
13+
count += 1;
14+
}
15+
assert_eq!(count, 3);
16+
}
17+
18+
fn main() {
19+
let future = real_main();
20+
let waker = noop_waker::NOOP_WAKER;
21+
let mut cx = &mut core::task::Context::from_waker(&waker);
22+
let mut future = core::pin::pin!(future);
23+
while let core::task::Poll::Pending = future.as_mut().poll(&mut cx) {}
24+
}
25+
26+
mod noop_waker {
27+
use std::task::{RawWaker, RawWakerVTable, Waker};
28+
29+
const VTABLE: RawWakerVTable =
30+
RawWakerVTable::new(|_| RawWaker::new(core::ptr::null(), &VTABLE), |_| (), |_| (), |_| ());
31+
32+
pub(super) const NOOP_WAKER: Waker =
33+
unsafe { Waker::from_raw(RawWaker::new(core::ptr::null(), &VTABLE)) };
34+
}

0 commit comments

Comments
 (0)