generator_light/
ext.rs

1use core::{ops::DerefMut, pin::Pin};
2
3use crate::{Generator, GeneratorState};
4
5pub struct Iter<P>(Pin<P>);
6
7impl<P> Iterator for Iter<P>
8where
9    P: DerefMut<Target: Generator<(), Return = ()>>,
10{
11    type Item = <P::Target as Generator<()>>::Yield;
12
13    fn next(&mut self) -> Option<Self::Item> {
14        let GeneratorState::Yield(val) = self.0.as_mut().resume(()) else {
15            return None;
16        };
17        Some(val)
18    }
19}
20
21pub fn from_iter<Item>(iter: impl IntoIterator<Item = Item>) -> impl Generator<Yield = Item> {
22    struct GenIter<I: Iterator>(I);
23    impl<I: Iterator> Generator for GenIter<I> {
24        type Return = ();
25        type Yield = I::Item;
26        fn resume(self: Pin<&mut Self>, _value: ()) -> GeneratorState<Self::Yield, Self::Return> {
27            // SAFETY: We are not moving out of the pinned field.
28            match unsafe { self.get_unchecked_mut() }.0.next() {
29                Some(val) => GeneratorState::Yield(val),
30                None => GeneratorState::Complete(()),
31            }
32        }
33    }
34    GenIter(iter.into_iter())
35}
36
37pub const fn from_fn<Resume, Yield, Return>(
38    f: impl FnMut(Resume) -> GeneratorState<Yield, Return>,
39) -> impl Generator<Resume, Yield = Yield, Return = Return> {
40    struct GenFn<F>(F);
41    impl<F, R, Y, Out> Generator<R> for GenFn<F>
42    where
43        F: FnMut(R) -> GeneratorState<Y, Out>,
44    {
45        type Return = Out;
46        type Yield = Y;
47        fn resume(self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
48            // SAFETY: We are not moving out of the pinned field.
49            unsafe { self.get_unchecked_mut() }.0(value)
50        }
51    }
52    GenFn(f)
53}
54
55pub trait GeneratorIterator {
56    type Item;
57    type Iter: Iterator<Item = Self::Item>;
58    fn into_iter(self) -> Self::Iter;
59}
60
61impl<P> GeneratorIterator for Pin<P>
62where
63    P: DerefMut<Target: Generator<(), Return = ()>>,
64{
65    type Item = <P::Target as Generator<()>>::Yield;
66    type Iter = Iter<P>;
67    fn into_iter(self) -> Self::Iter {
68        Iter(self)
69    }
70}
71
72#[derive(Debug, Clone, Copy)]
73pub enum Either<L, R> {
74    Left(L),
75    Right(R),
76}
77
78pub trait GeneratorExt<R>: Generator<R> {
79    fn compose<G: Generator<Self::Yield>>(self, then: G) -> Compose<Self, G>
80    where
81        Self: Sized,
82    {
83        Compose {
84            first: self,
85            second: then,
86        }
87    }
88
89    fn map_complete<U, F>(self, f: F) -> MapComplete<Self, F>
90    where
91        Self: Sized,
92        F: FnOnce(Self::Return) -> U,
93    {
94        MapComplete {
95            generator: self,
96            f: Some(f),
97        }
98    }
99}
100
101impl<G, R> GeneratorExt<R> for G where G: Generator<R> {}
102
103pub struct Compose<A, B> {
104    first: A,
105    second: B,
106}
107
108impl<A, B, R> Generator<R> for Compose<A, B>
109where
110    A: Generator<R>,
111    B: Generator<A::Yield>,
112{
113    type Return = Either<A::Return, B::Return>;
114    type Yield = B::Yield;
115
116    fn resume(mut self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
117        let first = unsafe { self.as_mut().map_unchecked_mut(|this| &mut this.first) };
118        let value = match first.resume(value) {
119            GeneratorState::Complete(x) => return GeneratorState::Complete(Either::Left(x)),
120            GeneratorState::Yield(value) => value,
121        };
122        let second = unsafe { self.map_unchecked_mut(|this| &mut this.second) };
123        second.resume(value).map_complete(Either::Right)
124    }
125}
126
127pub struct MapComplete<G, F> {
128    generator: G,
129    f: Option<F>,
130}
131
132impl<G, F, R, U> Generator<R> for MapComplete<G, F>
133where
134    G: Generator<R>,
135    F: FnOnce(G::Return) -> U,
136{
137    type Yield = G::Yield;
138    type Return = U;
139
140    fn resume(mut self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
141        let g = unsafe { self.as_mut().map_unchecked_mut(|this| &mut this.generator) };
142        g.resume(value)
143            .map_complete(|r| unsafe { self.get_unchecked_mut() }.f.take().unwrap()(r))
144    }
145}
generator_light/ext.rs

generator_light/
ext.rs
