Skip to content

Commit fbd39f9

Browse files
Implement async gen blocks
1 parent ec4ae32 commit fbd39f9

File tree

31 files changed

+549
-50
lines changed

31 files changed

+549
-50
lines changed

compiler/rustc_ast/src/ast.rs

+2
Original file line numberDiff line numberDiff line change
@@ -1501,6 +1501,7 @@ pub enum ExprKind {
15011501
pub enum GenBlockKind {
15021502
Async,
15031503
Gen,
1504+
AsyncGen,
15041505
}
15051506

15061507
impl fmt::Display for GenBlockKind {
@@ -1514,6 +1515,7 @@ impl GenBlockKind {
15141515
match self {
15151516
GenBlockKind::Async => "async",
15161517
GenBlockKind::Gen => "gen",
1518+
GenBlockKind::AsyncGen => "async gen",
15171519
}
15181520
}
15191521
}

compiler/rustc_ast_lowering/src/expr.rs

+125-13
Original file line numberDiff line numberDiff line change
@@ -325,6 +325,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
325325
hir::CoroutineSource::Block,
326326
|this| this.with_new_scopes(e.span, |this| this.lower_block_expr(block)),
327327
),
328+
ExprKind::Gen(capture_clause, block, GenBlockKind::AsyncGen) => self
329+
.make_async_gen_expr(
330+
*capture_clause,
331+
e.id,
332+
None,
333+
e.span,
334+
hir::CoroutineSource::Block,
335+
|this| this.with_new_scopes(e.span, |this| this.lower_block_expr(block)),
336+
),
328337
ExprKind::Yield(opt_expr) => self.lower_expr_yield(e.span, opt_expr.as_deref()),
329338
ExprKind::Err => hir::ExprKind::Err(
330339
self.tcx.sess.delay_span_bug(e.span, "lowered ExprKind::Err"),
@@ -726,6 +735,84 @@ impl<'hir> LoweringContext<'_, 'hir> {
726735
}))
727736
}
728737

738+
/// Lower a `async gen` construct to a generator that implements `AsyncIterator`.
739+
///
740+
/// This results in:
741+
///
742+
/// ```text
743+
/// static move? |_task_context| -> () {
744+
/// <body>
745+
/// }
746+
/// ```
747+
pub(super) fn make_async_gen_expr(
748+
&mut self,
749+
capture_clause: CaptureBy,
750+
closure_node_id: NodeId,
751+
_yield_ty: Option<hir::FnRetTy<'hir>>,
752+
span: Span,
753+
async_coroutine_source: hir::CoroutineSource,
754+
body: impl FnOnce(&mut Self) -> hir::Expr<'hir>,
755+
) -> hir::ExprKind<'hir> {
756+
let output = hir::FnRetTy::DefaultReturn(self.lower_span(span));
757+
758+
// Resume argument type: `ResumeTy`
759+
let unstable_span =
760+
self.mark_span_with_reason(DesugaringKind::Async, span, self.allow_gen_future.clone());
761+
let resume_ty = hir::QPath::LangItem(hir::LangItem::ResumeTy, unstable_span);
762+
let input_ty = hir::Ty {
763+
hir_id: self.next_id(),
764+
kind: hir::TyKind::Path(resume_ty),
765+
span: unstable_span,
766+
};
767+
768+
// The closure/coroutine `FnDecl` takes a single (resume) argument of type `input_ty`.
769+
let fn_decl = self.arena.alloc(hir::FnDecl {
770+
inputs: arena_vec![self; input_ty],
771+
output,
772+
c_variadic: false,
773+
implicit_self: hir::ImplicitSelfKind::None,
774+
lifetime_elision_allowed: false,
775+
});
776+
777+
// Lower the argument pattern/ident. The ident is used again in the `.await` lowering.
778+
let (pat, task_context_hid) = self.pat_ident_binding_mode(
779+
span,
780+
Ident::with_dummy_span(sym::_task_context),
781+
hir::BindingAnnotation::MUT,
782+
);
783+
let param = hir::Param {
784+
hir_id: self.next_id(),
785+
pat,
786+
ty_span: self.lower_span(span),
787+
span: self.lower_span(span),
788+
};
789+
let params = arena_vec![self; param];
790+
791+
let body = self.lower_body(move |this| {
792+
this.coroutine_kind = Some(hir::CoroutineKind::AsyncGen(async_coroutine_source));
793+
794+
let old_ctx = this.task_context;
795+
this.task_context = Some(task_context_hid);
796+
let res = body(this);
797+
this.task_context = old_ctx;
798+
(params, res)
799+
});
800+
801+
// `static |_task_context| -> <ret_ty> { body }`:
802+
hir::ExprKind::Closure(self.arena.alloc(hir::Closure {
803+
def_id: self.local_def_id(closure_node_id),
804+
binder: hir::ClosureBinder::Default,
805+
capture_clause,
806+
bound_generic_params: &[],
807+
fn_decl,
808+
body,
809+
fn_decl_span: self.lower_span(span),
810+
fn_arg_span: None,
811+
movability: Some(hir::Movability::Static),
812+
constness: hir::Constness::NotConst,
813+
}))
814+
}
815+
729816
/// Forwards a possible `#[track_caller]` annotation from `outer_hir_id` to
730817
/// `inner_hir_id` in case the `async_fn_track_caller` feature is enabled.
731818
pub(super) fn maybe_forward_track_caller(
@@ -775,15 +862,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
775862
/// ```
776863
fn lower_expr_await(&mut self, await_kw_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
777864
let full_span = expr.span.to(await_kw_span);
778-
match self.coroutine_kind {
779-
Some(hir::CoroutineKind::Async(_)) => {}
865+
866+
let is_async_gen = match self.coroutine_kind {
867+
Some(hir::CoroutineKind::Async(_)) => false,
868+
Some(hir::CoroutineKind::AsyncGen(_)) => true,
780869
Some(hir::CoroutineKind::Coroutine) | Some(hir::CoroutineKind::Gen(_)) | None => {
781870
return hir::ExprKind::Err(self.tcx.sess.emit_err(AwaitOnlyInAsyncFnAndBlocks {
782871
await_kw_span,
783872
item_span: self.current_item,
784873
}));
785874
}
786-
}
875+
};
876+
787877
let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, None);
788878
let gen_future_span = self.mark_span_with_reason(
789879
DesugaringKind::Await,
@@ -872,12 +962,19 @@ impl<'hir> LoweringContext<'_, 'hir> {
872962
self.stmt_expr(span, match_expr)
873963
};
874964

875-
// task_context = yield ();
965+
// Depending on `async` of `async gen`:
966+
// async - task_context = yield ();
967+
// async gen - task_context = yield async_gen_pending();
876968
let yield_stmt = {
877-
let unit = self.expr_unit(span);
969+
let yielded = if is_async_gen {
970+
self.expr_call_lang_item_fn(span, hir::LangItem::AsyncGenPending, &[])
971+
} else {
972+
self.expr_unit(span)
973+
};
974+
878975
let yield_expr = self.expr(
879976
span,
880-
hir::ExprKind::Yield(unit, hir::YieldSource::Await { expr: Some(expr_hir_id) }),
977+
hir::ExprKind::Yield(yielded, hir::YieldSource::Await { expr: Some(expr_hir_id) }),
881978
);
882979
let yield_expr = self.arena.alloc(yield_expr);
883980

@@ -987,7 +1084,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
9871084
}
9881085
Some(movability)
9891086
}
990-
Some(hir::CoroutineKind::Gen(_)) | Some(hir::CoroutineKind::Async(_)) => {
1087+
Some(
1088+
hir::CoroutineKind::Gen(_)
1089+
| hir::CoroutineKind::Async(_)
1090+
| hir::CoroutineKind::AsyncGen(_),
1091+
) => {
9911092
panic!("non-`async`/`gen` closure body turned `async`/`gen` during lowering");
9921093
}
9931094
None => {
@@ -1494,8 +1595,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
14941595
}
14951596

14961597
fn lower_expr_yield(&mut self, span: Span, opt_expr: Option<&Expr>) -> hir::ExprKind<'hir> {
1497-
match self.coroutine_kind {
1498-
Some(hir::CoroutineKind::Gen(_)) => {}
1598+
let is_async_gen = match self.coroutine_kind {
1599+
Some(hir::CoroutineKind::Gen(_)) => false,
1600+
Some(hir::CoroutineKind::AsyncGen(_)) => true,
14991601
Some(hir::CoroutineKind::Async(_)) => {
15001602
return hir::ExprKind::Err(
15011603
self.tcx.sess.emit_err(AsyncCoroutinesNotSupported { span }),
@@ -1511,14 +1613,24 @@ impl<'hir> LoweringContext<'_, 'hir> {
15111613
)
15121614
.emit();
15131615
}
1514-
self.coroutine_kind = Some(hir::CoroutineKind::Coroutine)
1616+
self.coroutine_kind = Some(hir::CoroutineKind::Coroutine);
1617+
false
15151618
}
1516-
}
1619+
};
15171620

1518-
let expr =
1621+
let mut yielded =
15191622
opt_expr.as_ref().map(|x| self.lower_expr(x)).unwrap_or_else(|| self.expr_unit(span));
15201623

1521-
hir::ExprKind::Yield(expr, hir::YieldSource::Yield)
1624+
if is_async_gen {
1625+
// yield async_gen_ready($expr);
1626+
yielded = self.expr_call_lang_item_fn(
1627+
span,
1628+
hir::LangItem::AsyncGenReady,
1629+
std::slice::from_ref(yielded),
1630+
);
1631+
}
1632+
1633+
hir::ExprKind::Yield(yielded, hir::YieldSource::Yield)
15221634
}
15231635

15241636
/// Desugar `ExprForLoop` from: `[opt_ident]: for <pat> in <head> <body>` into:

compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs

+5
Original file line numberDiff line numberDiff line change
@@ -2519,6 +2519,11 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
25192519
CoroutineSource::Closure => "gen closure",
25202520
_ => bug!("gen block/closure expected, but gen function found."),
25212521
},
2522+
CoroutineKind::AsyncGen(kind) => match kind {
2523+
CoroutineSource::Block => "async gen block",
2524+
CoroutineSource::Closure => "async gen closure",
2525+
_ => bug!("gen block/closure expected, but gen function found."),
2526+
},
25222527
CoroutineKind::Async(async_kind) => match async_kind {
25232528
CoroutineSource::Block => "async block",
25242529
CoroutineSource::Closure => "async closure",

compiler/rustc_borrowck/src/diagnostics/region_name.rs

+15
Original file line numberDiff line numberDiff line change
@@ -715,6 +715,21 @@ impl<'tcx> MirBorrowckCtxt<'_, 'tcx> {
715715
" of gen function"
716716
}
717717
},
718+
719+
Some(hir::CoroutineKind::AsyncGen(gen)) => match gen {
720+
hir::CoroutineSource::Block => " of async gen block",
721+
hir::CoroutineSource::Closure => " of async gen closure",
722+
hir::CoroutineSource::Fn => {
723+
let parent_item =
724+
hir.get_by_def_id(hir.get_parent_item(mir_hir_id).def_id);
725+
let output = &parent_item
726+
.fn_decl()
727+
.expect("coroutine lowered from async gen fn should be in fn")
728+
.output;
729+
span = output.span();
730+
" of async gen function"
731+
}
732+
},
718733
Some(hir::CoroutineKind::Coroutine) => " of coroutine",
719734
None => " of closure",
720735
};

compiler/rustc_codegen_ssa/src/debuginfo/type_names.rs

+3
Original file line numberDiff line numberDiff line change
@@ -566,6 +566,9 @@ fn coroutine_kind_label(coroutine_kind: Option<CoroutineKind>) -> &'static str {
566566
Some(CoroutineKind::Async(CoroutineSource::Block)) => "async_block",
567567
Some(CoroutineKind::Async(CoroutineSource::Closure)) => "async_closure",
568568
Some(CoroutineKind::Async(CoroutineSource::Fn)) => "async_fn",
569+
Some(CoroutineKind::AsyncGen(CoroutineSource::Block)) => "async_gen_block",
570+
Some(CoroutineKind::AsyncGen(CoroutineSource::Closure)) => "async_gen_closure",
571+
Some(CoroutineKind::AsyncGen(CoroutineSource::Fn)) => "async_gen_fn",
569572
Some(CoroutineKind::Coroutine) => "coroutine",
570573
None => "closure",
571574
}

compiler/rustc_codegen_ssa/src/mir/locals.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
4343
let local = mir::Local::from_usize(local);
4444
let expected_ty = self.monomorphize(self.mir.local_decls[local].ty);
4545
if expected_ty != op.layout.ty {
46-
warn!("Unexpected initial operand type. See the issues/114858");
46+
warn!(
47+
"Unexpected initial operand type: expected {expected_ty:?}, found {:?}.\
48+
See <https://github.com/rust-lang/rust/issues/114858>.",
49+
op.layout.ty
50+
);
4751
}
4852
}
4953
}

compiler/rustc_hir/src/hir.rs

+13-12
Original file line numberDiff line numberDiff line change
@@ -1484,12 +1484,16 @@ impl<'hir> Body<'hir> {
14841484
/// The type of source expression that caused this coroutine to be created.
14851485
#[derive(Clone, PartialEq, Eq, Debug, Copy, Hash, HashStable_Generic, Encodable, Decodable)]
14861486
pub enum CoroutineKind {
1487-
/// An explicit `async` block or the body of an async function.
1487+
/// An explicit `async` block or the body of an `async` function.
14881488
Async(CoroutineSource),
14891489

14901490
/// An explicit `gen` block or the body of a `gen` function.
14911491
Gen(CoroutineSource),
14921492

1493+
/// An explicit `async gen` block or the body of an `async gen` function,
1494+
/// which is able to both `yield` and `.await`.
1495+
AsyncGen(CoroutineSource),
1496+
14931497
/// A coroutine literal created via a `yield` inside a closure.
14941498
Coroutine,
14951499
}
@@ -1514,6 +1518,14 @@ impl fmt::Display for CoroutineKind {
15141518
}
15151519
k.fmt(f)
15161520
}
1521+
CoroutineKind::AsyncGen(k) => {
1522+
if f.alternate() {
1523+
f.write_str("`async gen` ")?;
1524+
} else {
1525+
f.write_str("async gen ")?
1526+
}
1527+
k.fmt(f)
1528+
}
15171529
}
15181530
}
15191531
}
@@ -2209,17 +2221,6 @@ impl fmt::Display for YieldSource {
22092221
}
22102222
}
22112223

2212-
impl From<CoroutineKind> for YieldSource {
2213-
fn from(kind: CoroutineKind) -> Self {
2214-
match kind {
2215-
// Guess based on the kind of the current coroutine.
2216-
CoroutineKind::Coroutine => Self::Yield,
2217-
CoroutineKind::Async(_) => Self::Await { expr: None },
2218-
CoroutineKind::Gen(_) => Self::Yield,
2219-
}
2220-
}
2221-
}
2222-
22232224
// N.B., if you change this, you'll probably want to change the corresponding
22242225
// type structure in middle/ty.rs as well.
22252226
#[derive(Debug, Clone, Copy, HashStable_Generic)]

compiler/rustc_hir/src/lang_items.rs

+5
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,7 @@ language_item_table! {
212212

213213
Iterator, sym::iterator, iterator_trait, Target::Trait, GenericRequirement::Exact(0);
214214
Future, sym::future_trait, future_trait, Target::Trait, GenericRequirement::Exact(0);
215+
AsyncIterator, sym::async_iterator, async_iterator_trait, Target::Trait, GenericRequirement::Exact(0);
215216
CoroutineState, sym::coroutine_state, coroutine_state, Target::Enum, GenericRequirement::None;
216217
Coroutine, sym::coroutine, coroutine_trait, Target::Trait, GenericRequirement::Minimum(1);
217218
Unpin, sym::unpin, unpin_trait, Target::Trait, GenericRequirement::None;
@@ -294,6 +295,10 @@ language_item_table! {
294295
PollReady, sym::Ready, poll_ready_variant, Target::Variant, GenericRequirement::None;
295296
PollPending, sym::Pending, poll_pending_variant, Target::Variant, GenericRequirement::None;
296297

298+
AsyncGenReady, sym::AsyncGenReady, async_gen_ready, Target::Method(MethodKind::Inherent), GenericRequirement::Exact(1);
299+
AsyncGenPending, sym::AsyncGenPending, async_gen_pending, Target::Method(MethodKind::Inherent), GenericRequirement::Exact(1);
300+
AsyncGenFinished, sym::AsyncGenFinished, async_gen_finished, Target::AssocConst, GenericRequirement::Exact(1);
301+
297302
// FIXME(swatinem): the following lang items are used for async lowering and
298303
// should become obsolete eventually.
299304
ResumeTy, sym::ResumeTy, resume_ty, Target::Struct, GenericRequirement::None;

compiler/rustc_hir_typeck/src/check.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,9 @@ pub(super) fn check_fn<'a, 'tcx>(
5959
&& can_be_coroutine.is_some()
6060
{
6161
let yield_ty = match kind {
62-
hir::CoroutineKind::Gen(..) | hir::CoroutineKind::Coroutine => {
62+
hir::CoroutineKind::Gen(..)
63+
| hir::CoroutineKind::AsyncGen(..)
64+
| hir::CoroutineKind::Coroutine => {
6365
let yield_ty = fcx.next_ty_var(TypeVariableOrigin {
6466
kind: TypeVariableOriginKind::TypeInference,
6567
span,

compiler/rustc_middle/src/mir/terminator.rs

+8
Original file line numberDiff line numberDiff line change
@@ -150,11 +150,17 @@ impl<O> AssertKind<O> {
150150
RemainderByZero(_) => "attempt to calculate the remainder with a divisor of zero",
151151
ResumedAfterReturn(CoroutineKind::Coroutine) => "coroutine resumed after completion",
152152
ResumedAfterReturn(CoroutineKind::Async(_)) => "`async fn` resumed after completion",
153+
ResumedAfterReturn(CoroutineKind::AsyncGen(_)) => {
154+
"`async gen fn` resumed after completion"
155+
}
153156
ResumedAfterReturn(CoroutineKind::Gen(_)) => {
154157
"`gen fn` should just keep returning `None` after completion"
155158
}
156159
ResumedAfterPanic(CoroutineKind::Coroutine) => "coroutine resumed after panicking",
157160
ResumedAfterPanic(CoroutineKind::Async(_)) => "`async fn` resumed after panicking",
161+
ResumedAfterPanic(CoroutineKind::AsyncGen(_)) => {
162+
"`async gen fn` resumed after panicking"
163+
}
158164
ResumedAfterPanic(CoroutineKind::Gen(_)) => {
159165
"`gen fn` should just keep returning `None` after panicking"
160166
}
@@ -245,13 +251,15 @@ impl<O> AssertKind<O> {
245251
DivisionByZero(_) => middle_assert_divide_by_zero,
246252
RemainderByZero(_) => middle_assert_remainder_by_zero,
247253
ResumedAfterReturn(CoroutineKind::Async(_)) => middle_assert_async_resume_after_return,
254+
ResumedAfterReturn(CoroutineKind::AsyncGen(_)) => todo!(),
248255
ResumedAfterReturn(CoroutineKind::Gen(_)) => {
249256
bug!("gen blocks can be resumed after they return and will keep returning `None`")
250257
}
251258
ResumedAfterReturn(CoroutineKind::Coroutine) => {
252259
middle_assert_coroutine_resume_after_return
253260
}
254261
ResumedAfterPanic(CoroutineKind::Async(_)) => middle_assert_async_resume_after_panic,
262+
ResumedAfterPanic(CoroutineKind::AsyncGen(_)) => todo!(),
255263
ResumedAfterPanic(CoroutineKind::Gen(_)) => middle_assert_gen_resume_after_panic,
256264
ResumedAfterPanic(CoroutineKind::Coroutine) => {
257265
middle_assert_coroutine_resume_after_panic

0 commit comments

Comments
 (0)