Skip to content

Commit 84b2f21

Browse files
asukaminato0721asukaminato0721
committed
fix
update based on comment
1 parent cc5bc48 commit 84b2f21

21 files changed

Lines changed: 621 additions & 121 deletions

File tree

‎crates/hir-def/src/expr_store/body.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,13 +74,15 @@ impl Body {
7474
let mut params = None;
7575

7676
let mut is_async_fn = false;
77+
let mut is_gen_fn = false;
7778
let InFile { file_id, value: body } = {
7879
match def {
7980
DefWithBodyId::FunctionId(f) => {
8081
let f = f.lookup(db);
8182
let src = f.source(db);
8283
params = src.value.param_list();
8384
is_async_fn = src.value.async_token().is_some();
85+
is_gen_fn = src.value.gen_token().is_some();
8486
src.map(|it| it.body().map(ast::Expr::from))
8587
}
8688
DefWithBodyId::ConstId(c) => {
@@ -101,7 +103,8 @@ impl Body {
101103
}
102104
};
103105
let module = def.module(db);
104-
let (body, source_map) = lower_body(db, def, file_id, module, params, body, is_async_fn);
106+
let (body, source_map) =
107+
lower_body(db, def, file_id, module, params, body, is_async_fn, is_gen_fn);
105108

106109
(Arc::new(body), source_map)
107110
}

‎crates/hir-def/src/expr_store/lower.rs‎

Lines changed: 162 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,8 @@ use crate::{
4646
},
4747
hir::{
4848
Array, Binding, BindingAnnotation, BindingId, BindingProblems, CaptureBy, ClosureKind,
49-
CoroutineSource, Expr, ExprId, Item, Label, LabelId, Literal, MatchArm, Movability,
50-
OffsetOf, Pat, PatId, RecordFieldPat, RecordLitField, RecordSpread, Statement,
49+
CoroutineKind, CoroutineSource, Expr, ExprId, Item, Label, LabelId, Literal, MatchArm,
50+
Movability, OffsetOf, Pat, PatId, RecordFieldPat, RecordLitField, RecordSpread, Statement,
5151
generics::GenericParams,
5252
},
5353
item_scope::BuiltinShadowMode,
@@ -71,6 +71,7 @@ pub(super) fn lower_body(
7171
parameters: Option<ast::ParamList>,
7272
body: Option<ast::Expr>,
7373
is_async_fn: bool,
74+
is_gen_fn: bool,
7475
) -> (Body, BodySourceMap) {
7576
// We cannot leave the root span map empty and let any identifier from it be treated as root,
7677
// because when inside nested macros `SyntaxContextId`s from the outer macro will be interleaved
@@ -175,6 +176,8 @@ pub(super) fn lower_body(
175176
DefWithBodyId::VariantId(..) => Awaitable::No("enum variant"),
176177
}
177178
},
179+
is_async_fn,
180+
is_gen_fn,
178181
);
179182
collector.store.inference_roots = Some(smallvec![(body_expr, RootExprOrigin::BodyRoot)]);
180183

@@ -375,12 +378,20 @@ pub(crate) fn lower_function(
375378
expr_collector.lower_type_ref_opt(ret_type.ty(), &mut ExprCollector::impl_trait_allocator)
376379
});
377380

378-
let return_type = if fn_.value.async_token().is_some() {
379-
let path = hir_expand::mod_path::path![core::future::Future];
381+
let return_type = if fn_.value.async_token().is_some() || fn_.value.gen_token().is_some() {
382+
let (path, assoc_name) =
383+
match (fn_.value.async_token().is_some(), fn_.value.gen_token().is_some()) {
384+
(true, true) => {
385+
(hir_expand::mod_path::path![core::async_iter::AsyncIterator], sym::Item)
386+
}
387+
(true, false) => (hir_expand::mod_path::path![core::future::Future], sym::Output),
388+
(false, true) => (hir_expand::mod_path::path![core::iter::Iterator], sym::Item),
389+
(false, false) => unreachable!(),
390+
};
380391
let mut generic_args: Vec<_> =
381392
std::iter::repeat_n(None, path.segments().len() - 1).collect();
382393
let binding = AssociatedTypeBinding {
383-
name: Name::new_symbol_root(sym::Output),
394+
name: Name::new_symbol_root(assoc_name),
384395
args: None,
385396
type_ref: Some(
386397
return_type
@@ -949,10 +960,11 @@ impl<'db> ExprCollector<'db> {
949960
/// into the body. This is to make sure that the future actually owns the
950961
/// arguments that are passed to the function, and to ensure things like
951962
/// drop order are stable.
952-
fn lower_async_block_with_moved_arguments(
963+
fn lower_coroutine_with_moved_arguments(
953964
&mut self,
954965
params: &mut [PatId],
955966
body: ExprId,
967+
kind: CoroutineKind,
956968
coroutine_source: CoroutineSource,
957969
) -> ExprId {
958970
let mut statements = Vec::new();
@@ -988,24 +1000,81 @@ impl<'db> ExprCollector<'db> {
9881000
*param = pat_id;
9891001
}
9901002

991-
let async_ = self.async_block(
992-
coroutine_source,
993-
// The default capture mode here is by-ref. Later on during upvar analysis,
994-
// we will force the captured arguments to by-move, but for async closures,
995-
// we want to make sure that we avoid unnecessarily moving captures, or else
996-
// all async closures would default to `FnOnce` as their calling mode.
997-
CaptureBy::Ref,
1003+
let capture_by = match kind {
1004+
// Async closures use by-ref by default so they don't all collapse to `FnOnce`.
1005+
CoroutineKind::Async => CaptureBy::Ref,
1006+
CoroutineKind::Gen | CoroutineKind::AsyncGen => CaptureBy::Value,
1007+
};
1008+
let coroutine = self.desugared_coroutine_expr(
1009+
ClosureKind::Coroutine { kind, source: coroutine_source },
1010+
capture_by,
9981011
None,
9991012
statements.into_boxed_slice(),
10001013
Some(body),
10011014
);
10021015
// It's important that this comes last, see the lowering of async closures for why.
1003-
self.alloc_expr_desugared(async_)
1016+
self.alloc_expr_desugared(coroutine)
1017+
}
1018+
1019+
/// Coroutine-like functions need to capture all parameters in the generated block, even if they have
1020+
/// non-captured patterns such as wildcards (to ensure consistent drop order).
1021+
fn lower_coroutine_fn(
1022+
&mut self,
1023+
params: &mut [PatId],
1024+
body: ExprId,
1025+
is_async_fn: bool,
1026+
is_gen_fn: bool,
1027+
) -> ExprId {
1028+
let mut statements = Vec::new();
1029+
for param in params {
1030+
let (name, hygiene) = match self.store.pats[*param] {
1031+
Pat::Bind { id, .. }
1032+
if matches!(
1033+
self.store.bindings[id].mode,
1034+
BindingAnnotation::Unannotated | BindingAnnotation::Mutable
1035+
) =>
1036+
{
1037+
continue;
1038+
}
1039+
Pat::Bind { id, .. } => {
1040+
(self.store.bindings[id].name.clone(), self.store.bindings[id].hygiene)
1041+
}
1042+
_ => (self.generate_new_name(), HygieneId::ROOT),
1043+
};
1044+
let binding_id = self.alloc_binding(name.clone(), BindingAnnotation::Mutable, hygiene);
1045+
let pat_id = self.alloc_pat_desugared(Pat::Bind { id: binding_id, subpat: None });
1046+
let expr = self.alloc_expr_desugared(Expr::Path(name.into()));
1047+
if !hygiene.is_root() {
1048+
self.store.ident_hygiene.insert(expr.into(), hygiene);
1049+
}
1050+
statements.push(Statement::Let {
1051+
pat: *param,
1052+
type_ref: None,
1053+
initializer: Some(expr),
1054+
else_branch: None,
1055+
});
1056+
*param = pat_id;
1057+
}
1058+
1059+
let kind = match (is_async_fn, is_gen_fn) {
1060+
(true, true) => CoroutineKind::AsyncGen,
1061+
(true, false) => CoroutineKind::Async,
1062+
(false, true) => CoroutineKind::Gen,
1063+
(false, false) => unreachable!(),
1064+
};
1065+
let coroutine = self.desugared_coroutine_expr(
1066+
ClosureKind::Coroutine { kind, source: CoroutineSource::Fn },
1067+
CaptureBy::Value,
1068+
None,
1069+
statements.into_boxed_slice(),
1070+
Some(body),
1071+
);
1072+
self.alloc_expr_desugared(coroutine)
10041073
}
10051074

1006-
fn async_block(
1075+
fn desugared_coroutine_expr(
10071076
&mut self,
1008-
source: CoroutineSource,
1077+
closure_kind: ClosureKind,
10091078
capture_by: CaptureBy,
10101079
id: Option<BlockId>,
10111080
statements: Box<[Statement]>,
@@ -1017,7 +1086,7 @@ impl<'db> ExprCollector<'db> {
10171086
arg_types: Box::default(),
10181087
ret_type: None,
10191088
body: block,
1020-
closure_kind: ClosureKind::AsyncBlock { source },
1089+
closure_kind,
10211090
capture_by,
10221091
}
10231092
}
@@ -1027,12 +1096,21 @@ impl<'db> ExprCollector<'db> {
10271096
params: &mut [PatId],
10281097
expr: Option<ast::Expr>,
10291098
awaitable: Awaitable,
1099+
is_async_fn: bool,
1100+
is_gen_fn: bool,
10301101
) -> ExprId {
10311102
self.awaitable_context.replace(awaitable);
10321103
self.with_label_rib(RibKind::Closure, |this| {
10331104
let body = this.collect_expr_opt(expr);
1034-
if awaitable == Awaitable::Yes {
1035-
this.lower_async_block_with_moved_arguments(params, body, CoroutineSource::Fn)
1105+
if is_gen_fn {
1106+
this.lower_coroutine_fn(params, body, is_async_fn, is_gen_fn)
1107+
} else if awaitable == Awaitable::Yes {
1108+
this.lower_coroutine_with_moved_arguments(
1109+
params,
1110+
body,
1111+
CoroutineKind::Async,
1112+
CoroutineSource::Fn,
1113+
)
10361114
} else {
10371115
body
10381116
}
@@ -1191,8 +1269,51 @@ impl<'db> ExprCollector<'db> {
11911269
self.with_label_rib(RibKind::Closure, |this| {
11921270
this.with_awaitable_block(Awaitable::Yes, |this| {
11931271
this.collect_block_(e, |this, id, statements, tail| {
1194-
this.async_block(
1195-
CoroutineSource::Block,
1272+
this.desugared_coroutine_expr(
1273+
ClosureKind::Coroutine {
1274+
kind: CoroutineKind::Async,
1275+
source: CoroutineSource::Block,
1276+
},
1277+
capture_by,
1278+
id,
1279+
statements,
1280+
tail,
1281+
)
1282+
})
1283+
})
1284+
})
1285+
}
1286+
Some(ast::BlockModifier::Gen(_)) => {
1287+
let capture_by =
1288+
if e.move_token().is_some() { CaptureBy::Value } else { CaptureBy::Ref };
1289+
self.with_label_rib(RibKind::Closure, |this| {
1290+
this.with_awaitable_block(Awaitable::No("non-async gen block"), |this| {
1291+
this.collect_block_(e, |this, id, statements, tail| {
1292+
this.desugared_coroutine_expr(
1293+
ClosureKind::Coroutine {
1294+
kind: CoroutineKind::Gen,
1295+
source: CoroutineSource::Block,
1296+
},
1297+
capture_by,
1298+
id,
1299+
statements,
1300+
tail,
1301+
)
1302+
})
1303+
})
1304+
})
1305+
}
1306+
Some(ast::BlockModifier::AsyncGen(_)) => {
1307+
let capture_by =
1308+
if e.move_token().is_some() { CaptureBy::Value } else { CaptureBy::Ref };
1309+
self.with_label_rib(RibKind::Closure, |this| {
1310+
this.with_awaitable_block(Awaitable::Yes, |this| {
1311+
this.collect_block_(e, |this, id, statements, tail| {
1312+
this.desugared_coroutine_expr(
1313+
ClosureKind::Coroutine {
1314+
kind: CoroutineKind::AsyncGen,
1315+
source: CoroutineSource::Block,
1316+
},
11961317
capture_by,
11971318
id,
11981319
statements,
@@ -1212,14 +1333,6 @@ impl<'db> ExprCollector<'db> {
12121333
})
12131334
})
12141335
}
1215-
// FIXME
1216-
Some(ast::BlockModifier::AsyncGen(_)) => {
1217-
self.with_awaitable_block(Awaitable::Yes, |this| this.collect_block(e))
1218-
}
1219-
Some(ast::BlockModifier::Gen(_)) => self
1220-
.with_awaitable_block(Awaitable::No("non-async gen block"), |this| {
1221-
this.collect_block(e)
1222-
}),
12231336
None => self.collect_block(e),
12241337
},
12251338
ast::Expr::LoopExpr(e) => {
@@ -1458,24 +1571,35 @@ impl<'db> ExprCollector<'db> {
14581571
let mut body = this
14591572
.with_awaitable_block(awaitable, |this| this.collect_expr_opt(e.body()));
14601573

1461-
let closure_kind = if this.is_lowering_coroutine {
1462-
let movability = if e.static_token().is_some() {
1463-
Movability::Static
1464-
} else {
1465-
Movability::Movable
1466-
};
1467-
ClosureKind::Coroutine(movability)
1574+
let closure_kind = if let Some(kind) =
1575+
if e.async_token().is_some() && e.gen_token().is_some()
1576+
{
1577+
Some(CoroutineKind::AsyncGen)
14681578
} else if e.async_token().is_some() {
1579+
Some(CoroutineKind::Async)
1580+
} else if e.gen_token().is_some() {
1581+
Some(CoroutineKind::Gen)
1582+
} else {
1583+
None
1584+
} {
14691585
// It's important that this expr is allocated immediately before the closure.
14701586
// We rely on it for `coroutine_for_closure()`.
1471-
body = this.lower_async_block_with_moved_arguments(
1587+
body = this.lower_coroutine_with_moved_arguments(
14721588
&mut args,
14731589
body,
1590+
kind,
14741591
CoroutineSource::Closure,
14751592
);
14761593
body_is_bindings_owner = true;
14771594

1478-
ClosureKind::AsyncClosure
1595+
ClosureKind::CoroutineClosure(kind)
1596+
} else if this.is_lowering_coroutine {
1597+
let movability = if e.static_token().is_some() {
1598+
Movability::Static
1599+
} else {
1600+
Movability::Movable
1601+
};
1602+
ClosureKind::OldCoroutine(movability)
14791603
} else {
14801604
ClosureKind::Closure
14811605
};

‎crates/hir-def/src/expr_store/pretty.rs‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ use crate::{
1717
attrs::AttrFlags,
1818
expr_store::path::{GenericArg, GenericArgs},
1919
hir::{
20-
Array, BindingAnnotation, CaptureBy, ClosureKind, Literal, Movability, RecordSpread,
21-
Statement,
20+
Array, BindingAnnotation, CaptureBy, ClosureKind, CoroutineKind, Literal, Movability,
21+
RecordSpread, Statement,
2222
generics::{GenericParams, WherePredicate},
2323
},
2424
lang_item::LangItemTarget,
@@ -764,28 +764,36 @@ impl Printer<'_> {
764764
let mut body = *body;
765765
let mut print_pipes = true;
766766
match closure_kind {
767-
ClosureKind::Coroutine(Movability::Static) => {
767+
ClosureKind::OldCoroutine(Movability::Static) => {
768768
w!(self, "static ");
769769
}
770-
ClosureKind::AsyncClosure => {
770+
ClosureKind::CoroutineClosure(kind) => {
771771
if let Expr::Closure {
772772
body: inner_body,
773-
closure_kind: ClosureKind::AsyncBlock { .. },
773+
closure_kind: ClosureKind::Coroutine { .. },
774774
..
775775
} = self.store[body]
776776
{
777777
body = inner_body;
778778
} else {
779-
never!("async closure should always have an async block body");
779+
never!("coroutine closure should always have a coroutine body");
780780
}
781781

782-
w!(self, "async ");
782+
match kind {
783+
CoroutineKind::Async => w!(self, "async "),
784+
CoroutineKind::Gen => w!(self, "gen "),
785+
CoroutineKind::AsyncGen => w!(self, "async gen "),
786+
}
783787
}
784-
ClosureKind::AsyncBlock { .. } => {
785-
w!(self, "async ");
788+
ClosureKind::Coroutine { kind, .. } => {
789+
match kind {
790+
CoroutineKind::Async => w!(self, "async "),
791+
CoroutineKind::Gen => w!(self, "gen "),
792+
CoroutineKind::AsyncGen => w!(self, "async gen "),
793+
}
786794
print_pipes = false;
787795
}
788-
ClosureKind::Closure | ClosureKind::Coroutine(Movability::Movable) => (),
796+
ClosureKind::Closure | ClosureKind::OldCoroutine(Movability::Movable) => (),
789797
}
790798
match capture_by {
791799
CaptureBy::Value => {

0 commit comments

Comments
 (0)