Skip to content

Commit

Permalink
Implement FusedIterator for gen block
Browse files Browse the repository at this point in the history
  • Loading branch information
ShoyuVanilla committed Mar 21, 2024
1 parent 03994e4 commit ae4c5c8
Show file tree
Hide file tree
Showing 10 changed files with 105 additions and 7 deletions.
1 change: 1 addition & 0 deletions compiler/rustc_hir/src/lang_items.rs
Expand Up @@ -214,6 +214,7 @@ language_item_table! {
FnOnceOutput, sym::fn_once_output, fn_once_output, Target::AssocTy, GenericRequirement::None;

Iterator, sym::iterator, iterator_trait, Target::Trait, GenericRequirement::Exact(0);
FusedIterator, sym::fused_iterator, fused_iterator_trait, Target::Trait, GenericRequirement::Exact(0);
Future, sym::future_trait, future_trait, Target::Trait, GenericRequirement::Exact(0);
AsyncIterator, sym::async_iterator, async_iterator_trait, Target::Trait, GenericRequirement::Exact(0);

Expand Down
2 changes: 2 additions & 0 deletions compiler/rustc_span/src/symbol.rs
Expand Up @@ -207,6 +207,7 @@ symbols! {
FromResidual,
FsOpenOptions,
FsPermissions,
FusedIterator,
Future,
FutureOutput,
GlobalAlloc,
Expand Down Expand Up @@ -885,6 +886,7 @@ symbols! {
fsub_algebraic,
fsub_fast,
fundamental,
fused_iterator,
future,
future_trait,
gdb_script_file,
Expand Down
9 changes: 9 additions & 0 deletions compiler/rustc_trait_selection/src/solve/assembly/mod.rs
Expand Up @@ -215,6 +215,13 @@ pub(super) trait GoalKind<'tcx>:
goal: Goal<'tcx, Self>,
) -> QueryResult<'tcx>;

/// A coroutine (that comes from a `gen` desugaring) is known to implement
/// `FusedIterator`
fn consider_builtin_fused_iterator_candidate(
ecx: &mut EvalCtxt<'_, 'tcx>,
goal: Goal<'tcx, Self>,
) -> QueryResult<'tcx>;

fn consider_builtin_async_iterator_candidate(
ecx: &mut EvalCtxt<'_, 'tcx>,
goal: Goal<'tcx, Self>,
Expand Down Expand Up @@ -497,6 +504,8 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
G::consider_builtin_future_candidate(self, goal)
} else if lang_items.iterator_trait() == Some(trait_def_id) {
G::consider_builtin_iterator_candidate(self, goal)
} else if lang_items.fused_iterator_trait() == Some(trait_def_id) {
G::consider_builtin_fused_iterator_candidate(self, goal)
} else if lang_items.async_iterator_trait() == Some(trait_def_id) {
G::consider_builtin_async_iterator_candidate(self, goal)
} else if lang_items.coroutine_trait() == Some(trait_def_id) {
Expand Down
7 changes: 7 additions & 0 deletions compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs
Expand Up @@ -647,6 +647,13 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
)
}

fn consider_builtin_fused_iterator_candidate(
_ecx: &mut EvalCtxt<'_, 'tcx>,
goal: Goal<'tcx, Self>,
) -> QueryResult<'tcx> {
bug!("`FusedIterator` does not have an associated type: {:?}", goal);
}

fn consider_builtin_async_iterator_candidate(
ecx: &mut EvalCtxt<'_, 'tcx>,
goal: Goal<'tcx, Self>,
Expand Down
22 changes: 22 additions & 0 deletions compiler/rustc_trait_selection/src/solve/trait_goals.rs
Expand Up @@ -456,6 +456,28 @@ impl<'tcx> assembly::GoalKind<'tcx> for TraitPredicate<'tcx> {
ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
}

fn consider_builtin_fused_iterator_candidate(
ecx: &mut EvalCtxt<'_, 'tcx>,
goal: Goal<'tcx, Self>,
) -> QueryResult<'tcx> {
if goal.predicate.polarity != ty::ImplPolarity::Positive {
return Err(NoSolution);
}

let ty::Coroutine(def_id, _) = *goal.predicate.self_ty().kind() else {
return Err(NoSolution);
};

// Coroutines are not iterators unless they come from `gen` desugaring
let tcx = ecx.tcx();
if !tcx.coroutine_is_gen(def_id) {
return Err(NoSolution);
}

// Gen coroutines unconditionally implement `FusedIterator`
ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
}

fn consider_builtin_async_iterator_candidate(
ecx: &mut EvalCtxt<'_, 'tcx>,
goal: Goal<'tcx, Self>,
Expand Down
Expand Up @@ -118,6 +118,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.fused_iterator_trait() == Some(def_id) {
self.assemble_fused_iterator_candidates(obligation, &mut candidates);
} else if lang_items.async_iterator_trait() == Some(def_id) {
self.assemble_async_iterator_candidates(obligation, &mut candidates);
} else if lang_items.async_fn_kind_helper() == Some(def_id) {
Expand Down Expand Up @@ -302,14 +304,31 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
candidates: &mut SelectionCandidateSet<'tcx>,
) {
let self_ty = obligation.self_ty().skip_binder();
if let ty::Coroutine(did, ..) = self_ty.kind() {
// gen constructs get lowered to a special kind of coroutine that
// should directly `impl Iterator`.
if self.tcx().coroutine_is_gen(*did) {
debug!(?self_ty, ?obligation, "assemble_iterator_candidates",);
// gen constructs get lowered to a special kind of coroutine that
// should directly `impl Iterator`.
if let ty::Coroutine(did, ..) = self_ty.kind()
&& self.tcx().coroutine_is_gen(*did)
{
debug!(?self_ty, ?obligation, "assemble_iterator_candidates",);

candidates.vec.push(IteratorCandidate);
}
candidates.vec.push(IteratorCandidate);
}
}

fn assemble_fused_iterator_candidates(
&mut self,
obligation: &PolyTraitObligation<'tcx>,
candidates: &mut SelectionCandidateSet<'tcx>,
) {
let self_ty = obligation.self_ty().skip_binder();
// gen constructs get lowered to a special kind of coroutine that
// should directly `impl FusedIterator`.
if let ty::Coroutine(did, ..) = self_ty.kind()
&& self.tcx().coroutine_is_gen(*did)
{
debug!(?self_ty, ?obligation, "assemble_fused_iterator_candidates",);

candidates.vec.push(BuiltinCandidate { has_nested: false });
}
}

Expand Down
Expand Up @@ -267,6 +267,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
self.copy_clone_conditions(obligation)
} else if Some(trait_def) == lang_items.clone_trait() {
self.copy_clone_conditions(obligation)
} else if Some(trait_def) == lang_items.fused_iterator_trait() {
self.fused_iterator_conditions(obligation)
} else {
bug!("unexpected builtin trait {:?}", trait_def)
};
Expand Down
14 changes: 14 additions & 0 deletions compiler/rustc_trait_selection/src/traits/select/mod.rs
Expand Up @@ -2259,6 +2259,20 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
}
}

fn fused_iterator_conditions(
&mut self,
obligation: &PolyTraitObligation<'tcx>,
) -> BuiltinImplConditions<'tcx> {
let self_ty = self.infcx.shallow_resolve(obligation.self_ty().skip_binder());
if let ty::Coroutine(did, ..) = *self_ty.kind()
&& self.tcx().coroutine_is_gen(did)
{
BuiltinImplConditions::Where(ty::Binder::dummy(Vec::new()))
} else {
BuiltinImplConditions::None
}
}

/// For default impls, we need to break apart a type into its
/// "constituent types" -- meaning, the types that it contains.
///
Expand Down
1 change: 1 addition & 0 deletions library/core/src/iter/traits/marker.rs
Expand Up @@ -28,6 +28,7 @@ pub unsafe trait TrustedFused {}
#[rustc_unsafe_specialization_marker]
// FIXME: this should be a #[marker] and have another blanket impl for T: TrustedFused
// but that ICEs iter::Fuse specializations.
#[cfg_attr(not(bootstrap), lang = "fused_iterator")]
pub trait FusedIterator: Iterator {}

#[stable(feature = "fused", since = "1.26.0")]
Expand Down
21 changes: 21 additions & 0 deletions tests/ui/coroutine/gen_block_is_fused_iter.rs
@@ -0,0 +1,21 @@
//@ revisions: next old
//@compile-flags: --edition 2024 -Zunstable-options
//@[next] compile-flags: -Znext-solver
//@ check-pass
#![feature(gen_blocks)]

use std::iter::FusedIterator;

fn foo() -> impl FusedIterator {
gen { yield 42 }
}

fn bar() -> impl FusedIterator<Item = u16> {
gen { yield 42 }
}

fn baz() -> impl FusedIterator + Iterator<Item = i64> {
gen { yield 42 }
}

fn main() {}

0 comments on commit ae4c5c8

Please sign in to comment.