Implement async gen blocks

This commit is contained in:
Michael Goulet
2023-11-28 18:18:19 +00:00
parent a0cbc168c9
commit 96bb542a31
32 changed files with 563 additions and 54 deletions

View File

@@ -2587,6 +2587,23 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
CoroutineKind::Async(CoroutineSource::Closure) => {
format!("future created by async closure is not {trait_name}")
}
CoroutineKind::AsyncGen(CoroutineSource::Fn) => self
.tcx
.parent(coroutine_did)
.as_local()
.map(|parent_did| self.tcx.local_def_id_to_hir_id(parent_did))
.and_then(|parent_hir_id| hir.opt_name(parent_hir_id))
.map(|name| {
format!("async iterator returned by `{name}` is not {trait_name}")
})?,
CoroutineKind::AsyncGen(CoroutineSource::Block) => {
format!("async iterator created by async gen block is not {trait_name}")
}
CoroutineKind::AsyncGen(CoroutineSource::Closure) => {
format!(
"async iterator created by async gen closure is not {trait_name}"
)
}
CoroutineKind::Gen(CoroutineSource::Fn) => self
.tcx
.parent(coroutine_did)
@@ -3127,7 +3144,9 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
let what = match self.tcx.coroutine_kind(coroutine_def_id) {
None
| Some(hir::CoroutineKind::Coroutine)
| Some(hir::CoroutineKind::Gen(_)) => "yield",
| Some(hir::CoroutineKind::Gen(_))
// FIXME(gen_blocks): This could be yield or await...
| Some(hir::CoroutineKind::AsyncGen(_)) => "yield",
Some(hir::CoroutineKind::Async(..)) => "await",
};
err.note(format!(

View File

@@ -1921,6 +1921,9 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
hir::CoroutineKind::Async(hir::CoroutineSource::Block) => "an async block",
hir::CoroutineKind::Async(hir::CoroutineSource::Fn) => "an async function",
hir::CoroutineKind::Async(hir::CoroutineSource::Closure) => "an async closure",
hir::CoroutineKind::AsyncGen(hir::CoroutineSource::Block) => "an async gen block",
hir::CoroutineKind::AsyncGen(hir::CoroutineSource::Fn) => "an async gen function",
hir::CoroutineKind::AsyncGen(hir::CoroutineSource::Closure) => "an async gen closure",
hir::CoroutineKind::Gen(hir::CoroutineSource::Block) => "a gen block",
hir::CoroutineKind::Gen(hir::CoroutineSource::Fn) => "a gen function",
hir::CoroutineKind::Gen(hir::CoroutineSource::Closure) => "a gen closure",

View File

@@ -1823,11 +1823,18 @@ fn assemble_candidates_from_impls<'cx, 'tcx>(
let self_ty = selcx.infcx.shallow_resolve(obligation.predicate.self_ty());
let lang_items = selcx.tcx().lang_items();
if [lang_items.coroutine_trait(), lang_items.future_trait(), lang_items.iterator_trait()].contains(&Some(trait_ref.def_id))
|| selcx.tcx().fn_trait_kind_from_def_id(trait_ref.def_id).is_some()
if [
lang_items.coroutine_trait(),
lang_items.future_trait(),
lang_items.iterator_trait(),
lang_items.async_iterator_trait(),
lang_items.fn_trait(),
lang_items.fn_mut_trait(),
lang_items.fn_once_trait(),
].contains(&Some(trait_ref.def_id))
{
true
} else if lang_items.discriminant_kind_trait() == Some(trait_ref.def_id) {
}else if lang_items.discriminant_kind_trait() == Some(trait_ref.def_id) {
match self_ty.kind() {
ty::Bool
| ty::Char
@@ -2042,6 +2049,8 @@ fn confirm_select_candidate<'cx, 'tcx>(
confirm_future_candidate(selcx, obligation, data)
} else if lang_items.iterator_trait() == Some(trait_def_id) {
confirm_iterator_candidate(selcx, obligation, data)
} else if lang_items.async_iterator_trait() == Some(trait_def_id) {
confirm_async_iterator_candidate(selcx, obligation, data)
} else if selcx.tcx().fn_trait_kind_from_def_id(trait_def_id).is_some() {
if obligation.predicate.self_ty().is_closure() {
confirm_closure_candidate(selcx, obligation, data)
@@ -2203,6 +2212,57 @@ fn confirm_iterator_candidate<'cx, 'tcx>(
.with_addl_obligations(obligations)
}
fn confirm_async_iterator_candidate<'cx, 'tcx>(
selcx: &mut SelectionContext<'cx, 'tcx>,
obligation: &ProjectionTyObligation<'tcx>,
nested: Vec<PredicateObligation<'tcx>>,
) -> Progress<'tcx> {
let ty::Coroutine(_, args, _) =
selcx.infcx.shallow_resolve(obligation.predicate.self_ty()).kind()
else {
unreachable!()
};
let gen_sig = args.as_coroutine().sig();
let Normalized { value: gen_sig, obligations } = normalize_with_depth(
selcx,
obligation.param_env,
obligation.cause.clone(),
obligation.recursion_depth + 1,
gen_sig,
);
debug!(?obligation, ?gen_sig, ?obligations, "confirm_async_iterator_candidate");
let tcx = selcx.tcx();
let iter_def_id = tcx.require_lang_item(LangItem::AsyncIterator, None);
let (trait_ref, yield_ty) = super::util::async_iterator_trait_ref_and_outputs(
tcx,
iter_def_id,
obligation.predicate.self_ty(),
gen_sig,
);
debug_assert_eq!(tcx.associated_item(obligation.predicate.def_id).name, sym::Item);
let ty::Adt(_poll_adt, args) = *yield_ty.kind() else {
bug!();
};
let ty::Adt(_option_adt, args) = *args.type_at(0).kind() else {
bug!();
};
let item_ty = args.type_at(0);
let predicate = ty::ProjectionPredicate {
projection_ty: ty::AliasTy::new(tcx, obligation.predicate.def_id, trait_ref.args),
term: item_ty.into(),
};
confirm_param_env_candidate(selcx, obligation, ty::Binder::dummy(predicate), false)
.with_addl_obligations(nested)
.with_addl_obligations(obligations)
}
fn confirm_builtin_candidate<'cx, 'tcx>(
selcx: &mut SelectionContext<'cx, 'tcx>,
obligation: &ProjectionTyObligation<'tcx>,

View File

@@ -112,6 +112,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
self.assemble_future_candidates(obligation, &mut candidates);
} else if lang_items.iterator_trait() == Some(def_id) {
self.assemble_iterator_candidates(obligation, &mut candidates);
} else if lang_items.async_iterator_trait() == Some(def_id) {
self.assemble_async_iterator_candidates(obligation, &mut candidates);
}
self.assemble_closure_candidates(obligation, &mut candidates);
@@ -258,6 +260,34 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
}
}
fn assemble_async_iterator_candidates(
&mut self,
obligation: &PolyTraitObligation<'tcx>,
candidates: &mut SelectionCandidateSet<'tcx>,
) {
let self_ty = obligation.self_ty().skip_binder();
if let ty::Coroutine(did, args, _) = *self_ty.kind() {
// gen constructs get lowered to a special kind of coroutine that
// should directly `impl AsyncIterator`.
if self.tcx().coroutine_is_async_gen(did) {
debug!(?self_ty, ?obligation, "assemble_iterator_candidates",);
// Can only confirm this candidate if we have constrained
// the `Yield` type to at least `Poll<Option<?0>>`..
let ty::Adt(_poll_def, args) = *args.as_coroutine().yield_ty().kind() else {
candidates.ambiguous = true;
return;
};
let ty::Adt(_option_def, _) = *args.type_at(0).kind() else {
candidates.ambiguous = true;
return;
};
candidates.vec.push(AsyncIteratorCandidate);
}
}
}
/// Checks for the artificial impl that the compiler will create for an obligation like `X :
/// FnMut<..>` where `X` is a closure type.
///

View File

@@ -98,6 +98,11 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
ImplSource::Builtin(BuiltinImplSource::Misc, vtable_iterator)
}
AsyncIteratorCandidate => {
let vtable_iterator = self.confirm_async_iterator_candidate(obligation)?;
ImplSource::Builtin(BuiltinImplSource::Misc, vtable_iterator)
}
FnPointerCandidate { is_const } => {
let data = self.confirm_fn_pointer_candidate(obligation, is_const)?;
ImplSource::Builtin(BuiltinImplSource::Misc, data)
@@ -813,6 +818,35 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
Ok(nested)
}
fn confirm_async_iterator_candidate(
&mut self,
obligation: &PolyTraitObligation<'tcx>,
) -> Result<Vec<PredicateObligation<'tcx>>, SelectionError<'tcx>> {
// Okay to skip binder because the args on coroutine types never
// touch bound regions, they just capture the in-scope
// type/region parameters.
let self_ty = self.infcx.shallow_resolve(obligation.self_ty().skip_binder());
let ty::Coroutine(coroutine_def_id, args, _) = *self_ty.kind() else {
bug!("closure candidate for non-closure {:?}", obligation);
};
debug!(?obligation, ?coroutine_def_id, ?args, "confirm_async_iterator_candidate");
let gen_sig = args.as_coroutine().sig();
let (trait_ref, _) = super::util::async_iterator_trait_ref_and_outputs(
self.tcx(),
obligation.predicate.def_id(),
obligation.predicate.no_bound_vars().expect("iterator has no bound vars").self_ty(),
gen_sig,
);
let nested = self.confirm_poly_trait_refs(obligation, ty::Binder::dummy(trait_ref))?;
debug!(?trait_ref, ?nested, "iterator candidate obligations");
Ok(nested)
}
#[instrument(skip(self), level = "debug")]
fn confirm_closure_candidate(
&mut self,

View File

@@ -1875,6 +1875,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
| CoroutineCandidate
| FutureCandidate
| IteratorCandidate
| AsyncIteratorCandidate
| FnPointerCandidate { .. }
| BuiltinObjectCandidate
| BuiltinUnsizeCandidate
@@ -1904,6 +1905,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
| CoroutineCandidate
| FutureCandidate
| IteratorCandidate
| AsyncIteratorCandidate
| FnPointerCandidate { .. }
| BuiltinObjectCandidate
| BuiltinUnsizeCandidate
@@ -1939,6 +1941,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
| CoroutineCandidate
| FutureCandidate
| IteratorCandidate
| AsyncIteratorCandidate
| FnPointerCandidate { .. }
| BuiltinObjectCandidate
| BuiltinUnsizeCandidate
@@ -1954,6 +1957,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
| CoroutineCandidate
| FutureCandidate
| IteratorCandidate
| AsyncIteratorCandidate
| FnPointerCandidate { .. }
| BuiltinObjectCandidate
| BuiltinUnsizeCandidate
@@ -2061,6 +2065,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
| CoroutineCandidate
| FutureCandidate
| IteratorCandidate
| AsyncIteratorCandidate
| FnPointerCandidate { .. }
| BuiltinObjectCandidate
| BuiltinUnsizeCandidate
@@ -2072,6 +2077,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
| CoroutineCandidate
| FutureCandidate
| IteratorCandidate
| AsyncIteratorCandidate
| FnPointerCandidate { .. }
| BuiltinObjectCandidate
| BuiltinUnsizeCandidate

View File

@@ -308,6 +308,17 @@ pub fn iterator_trait_ref_and_outputs<'tcx>(
(trait_ref, sig.yield_ty)
}
pub fn async_iterator_trait_ref_and_outputs<'tcx>(
tcx: TyCtxt<'tcx>,
async_iterator_def_id: DefId,
self_ty: Ty<'tcx>,
sig: ty::GenSig<'tcx>,
) -> (ty::TraitRef<'tcx>, Ty<'tcx>) {
assert!(!self_ty.has_escaping_bound_vars());
let trait_ref = ty::TraitRef::new(tcx, async_iterator_def_id, [self_ty]);
(trait_ref, sig.yield_ty)
}
pub fn impl_item_is_final(tcx: TyCtxt<'_>, assoc_item: &ty::AssocItem) -> bool {
assoc_item.defaultness(tcx).is_final()
&& tcx.defaultness(assoc_item.container_id(tcx)).is_final()