Skip to content

Commit 0561287

Browse files
compiler-errorscompiler-errors
committed
Implement async gen blocks
1 parent 1c89658 commit 0561287

File tree

31 files changed

+553
-50
lines changed

31 files changed

+553
-50
lines changed

‎compiler/rustc_ast/src/ast.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1516,6 +1516,7 @@ pub enum ExprKind {
15161516
pub enum GenBlockKind {
15171517
Async,
15181518
Gen,
1519+
AsyncGen,
15191520
}
15201521

15211522
impl fmt::Display for GenBlockKind {
@@ -1529,6 +1530,7 @@ impl GenBlockKind {
15291530
match self {
15301531
GenBlockKind::Async => "async",
15311532
GenBlockKind::Gen => "gen",
1533+
GenBlockKind::AsyncGen => "async gen",
15321534
}
15331535
}
15341536
}

‎compiler/rustc_ast_lowering/src/expr.rs‎

Lines changed: 128 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
324324
hir::CoroutineSource::Block,
325325
|this| this.with_new_scopes(e.span, |this| this.lower_block_expr(block)),
326326
),
327+
ExprKind::Gen(capture_clause, block, GenBlockKind::AsyncGen) => self
328+
.make_async_gen_expr(
329+
*capture_clause,
330+
e.id,
331+
None,
332+
e.span,
333+
hir::CoroutineSource::Block,
334+
|this| this.with_new_scopes(e.span, |this| this.lower_block_expr(block)),
335+
),
327336
ExprKind::Yield(opt_expr) => self.lower_expr_yield(e.span, opt_expr.as_deref()),
328337
ExprKind::Err => hir::ExprKind::Err(
329338
self.tcx.sess.span_delayed_bug(e.span, "lowered ExprKind::Err"),
@@ -706,6 +715,87 @@ impl<'hir> LoweringContext<'_, 'hir> {
706715
}))
707716
}
708717

718+
/// Lower a `async gen` construct to a generator that implements `AsyncIterator`.
719+
///
720+
/// This results in:
721+
///
722+
/// ```text
723+
/// static move? |_task_context| -> () {
724+
/// <body>
725+
/// }
726+
/// ```
727+
pub(super) fn make_async_gen_expr(
728+
&mut self,
729+
capture_clause: CaptureBy,
730+
closure_node_id: NodeId,
731+
_yield_ty: Option<hir::FnRetTy<'hir>>,
732+
span: Span,
733+
async_coroutine_source: hir::CoroutineSource,
734+
body: impl FnOnce(&mut Self) -> hir::Expr<'hir>,
735+
) -> hir::ExprKind<'hir> {
736+
let output = hir::FnRetTy::DefaultReturn(self.lower_span(span));
737+
738+
// Resume argument type: `ResumeTy`
739+
let unstable_span = self.mark_span_with_reason(
740+
DesugaringKind::Async,
741+
span,
742+
Some(self.allow_gen_future.clone()),
743+
);
744+
let resume_ty = hir::QPath::LangItem(hir::LangItem::ResumeTy, unstable_span);
745+
let input_ty = hir::Ty {
746+
hir_id: self.next_id(),
747+
kind: hir::TyKind::Path(resume_ty),
748+
span: unstable_span,
749+
};
750+
751+
// The closure/coroutine `FnDecl` takes a single (resume) argument of type `input_ty`.
752+
let fn_decl = self.arena.alloc(hir::FnDecl {
753+
inputs: arena_vec![self; input_ty],
754+
output,
755+
c_variadic: false,
756+
implicit_self: hir::ImplicitSelfKind::None,
757+
lifetime_elision_allowed: false,
758+
});
759+
760+
// Lower the argument pattern/ident. The ident is used again in the `.await` lowering.
761+
let (pat, task_context_hid) = self.pat_ident_binding_mode(
762+
span,
763+
Ident::with_dummy_span(sym::_task_context),
764+
hir::BindingAnnotation::MUT,
765+
);
766+
let param = hir::Param {
767+
hir_id: self.next_id(),
768+
pat,
769+
ty_span: self.lower_span(span),
770+
span: self.lower_span(span),
771+
};
772+
let params = arena_vec![self; param];
773+
774+
let body = self.lower_body(move |this| {
775+
this.coroutine_kind = Some(hir::CoroutineKind::AsyncGen(async_coroutine_source));
776+
777+
let old_ctx = this.task_context;
778+
this.task_context = Some(task_context_hid);
779+
let res = body(this);
780+
this.task_context = old_ctx;
781+
(params, res)
782+
});
783+
784+
// `static |_task_context| -> <ret_ty> { body }`:
785+
hir::ExprKind::Closure(self.arena.alloc(hir::Closure {
786+
def_id: self.local_def_id(closure_node_id),
787+
binder: hir::ClosureBinder::Default,
788+
capture_clause,
789+
bound_generic_params: &[],
790+
fn_decl,
791+
body,
792+
fn_decl_span: self.lower_span(span),
793+
fn_arg_span: None,
794+
movability: Some(hir::Movability::Static),
795+
constness: hir::Constness::NotConst,
796+
}))
797+
}
798+
709799
/// Forwards a possible `#[track_caller]` annotation from `outer_hir_id` to
710800
/// `inner_hir_id` in case the `async_fn_track_caller` feature is enabled.
711801
pub(super) fn maybe_forward_track_caller(
@@ -755,15 +845,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
755845
/// ```
756846
fn lower_expr_await(&mut self, await_kw_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
757847
let full_span = expr.span.to(await_kw_span);
758-
match self.coroutine_kind {
759-
Some(hir::CoroutineKind::Async(_)) => {}
848+
849+
let is_async_gen = match self.coroutine_kind {
850+
Some(hir::CoroutineKind::Async(_)) => false,
851+
Some(hir::CoroutineKind::AsyncGen(_)) => true,
760852
Some(hir::CoroutineKind::Coroutine) | Some(hir::CoroutineKind::Gen(_)) | None => {
761853
return hir::ExprKind::Err(self.tcx.sess.emit_err(AwaitOnlyInAsyncFnAndBlocks {
762854
await_kw_span,
763855
item_span: self.current_item,
764856
}));
765857
}
766-
}
858+
};
859+
767860
let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, None);
768861
let gen_future_span = self.mark_span_with_reason(
769862
DesugaringKind::Await,
@@ -852,12 +945,19 @@ impl<'hir> LoweringContext<'_, 'hir> {
852945
self.stmt_expr(span, match_expr)
853946
};
854947

855-
// task_context = yield ();
948+
// Depending on `async` of `async gen`:
949+
// async - task_context = yield ();
950+
// async gen - task_context = yield async_gen_pending();
856951
let yield_stmt = {
857-
let unit = self.expr_unit(span);
952+
let yielded = if is_async_gen {
953+
self.expr_call_lang_item_fn(span, hir::LangItem::AsyncGenPending, &[])
954+
} else {
955+
self.expr_unit(span)
956+
};
957+
858958
let yield_expr = self.expr(
859959
span,
860-
hir::ExprKind::Yield(unit, hir::YieldSource::Await { expr: Some(expr_hir_id) }),
960+
hir::ExprKind::Yield(yielded, hir::YieldSource::Await { expr: Some(expr_hir_id) }),
861961
);
862962
let yield_expr = self.arena.alloc(yield_expr);
863963

@@ -967,7 +1067,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
9671067
}
9681068
Some(movability)
9691069
}
970-
Some(hir::CoroutineKind::Gen(_)) | Some(hir::CoroutineKind::Async(_)) => {
1070+
Some(
1071+
hir::CoroutineKind::Gen(_)
1072+
| hir::CoroutineKind::Async(_)
1073+
| hir::CoroutineKind::AsyncGen(_),
1074+
) => {
9711075
panic!("non-`async`/`gen` closure body turned `async`/`gen` during lowering");
9721076
}
9731077
None => {
@@ -1474,8 +1578,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
14741578
}
14751579

14761580
fn lower_expr_yield(&mut self, span: Span, opt_expr: Option<&Expr>) -> hir::ExprKind<'hir> {
1477-
match self.coroutine_kind {
1478-
Some(hir::CoroutineKind::Gen(_)) => {}
1581+
let is_async_gen = match self.coroutine_kind {
1582+
Some(hir::CoroutineKind::Gen(_)) => false,
1583+
Some(hir::CoroutineKind::AsyncGen(_)) => true,
14791584
Some(hir::CoroutineKind::Async(_)) => {
14801585
return hir::ExprKind::Err(
14811586
self.tcx.sess.emit_err(AsyncCoroutinesNotSupported { span }),
@@ -1491,14 +1596,24 @@ impl<'hir> LoweringContext<'_, 'hir> {
14911596
)
14921597
.emit();
14931598
}
1494-
self.coroutine_kind = Some(hir::CoroutineKind::Coroutine)
1599+
self.coroutine_kind = Some(hir::CoroutineKind::Coroutine);
1600+
false
14951601
}
1496-
}
1602+
};
14971603

1498-
let expr =
1604+
let mut yielded =
14991605
opt_expr.as_ref().map(|x| self.lower_expr(x)).unwrap_or_else(|| self.expr_unit(span));
15001606

1501-
hir::ExprKind::Yield(expr, hir::YieldSource::Yield)
1607+
if is_async_gen {
1608+
// yield async_gen_ready($expr);
1609+
yielded = self.expr_call_lang_item_fn(
1610+
span,
1611+
hir::LangItem::AsyncGenReady,
1612+
std::slice::from_ref(yielded),
1613+
);
1614+
}
1615+
1616+
hir::ExprKind::Yield(yielded, hir::YieldSource::Yield)
15021617
}
15031618

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

‎compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs‎

Lines changed: 5 additions & 0 deletions
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‎

Lines changed: 15 additions & 0 deletions
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‎

Lines changed: 3 additions & 0 deletions
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‎

Lines changed: 5 additions & 1 deletion
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‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1339,12 +1339,16 @@ impl<'hir> Body<'hir> {
13391339
/// The type of source expression that caused this coroutine to be created.
13401340
#[derive(Clone, PartialEq, Eq, Debug, Copy, Hash, HashStable_Generic, Encodable, Decodable)]
13411341
pub enum CoroutineKind {
1342-
/// An explicit `async` block or the body of an async function.
1342+
/// An explicit `async` block or the body of an `async` function.
13431343
Async(CoroutineSource),
13441344

13451345
/// An explicit `gen` block or the body of a `gen` function.
13461346
Gen(CoroutineSource),
13471347

1348+
/// An explicit `async gen` block or the body of an `async gen` function,
1349+
/// which is able to both `yield` and `.await`.
1350+
AsyncGen(CoroutineSource),
1351+
13481352
/// A coroutine literal created via a `yield` inside a closure.
13491353
Coroutine,
13501354
}
@@ -1369,6 +1373,14 @@ impl fmt::Display for CoroutineKind {
13691373
}
13701374
k.fmt(f)
13711375
}
1376+
CoroutineKind::AsyncGen(k) => {
1377+
if f.alternate() {
1378+
f.write_str("`async gen` ")?;
1379+
} else {
1380+
f.write_str("async gen ")?
1381+
}
1382+
k.fmt(f)
1383+
}
13721384
}
13731385
}
13741386
}
@@ -2064,17 +2076,6 @@ impl fmt::Display for YieldSource {
20642076
}
20652077
}
20662078

2067-
impl From<CoroutineKind> for YieldSource {
2068-
fn from(kind: CoroutineKind) -> Self {
2069-
match kind {
2070-
// Guess based on the kind of the current coroutine.
2071-
CoroutineKind::Coroutine => Self::Yield,
2072-
CoroutineKind::Async(_) => Self::Await { expr: None },
2073-
CoroutineKind::Gen(_) => Self::Yield,
2074-
}
2075-
}
2076-
}
2077-
20782079
// N.B., if you change this, you'll probably want to change the corresponding
20792080
// type structure in middle/ty.rs as well.
20802081
#[derive(Debug, Clone, Copy, HashStable_Generic)]

‎compiler/rustc_hir/src/lang_items.rs‎

Lines changed: 5 additions & 0 deletions
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‎

Lines changed: 3 additions & 1 deletion
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‎

Lines changed: 8 additions & 0 deletions
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)