generator_light/
ext.rs

1use std::marker::PhantomData;
2
3use crate::core::convert::Infallible;
4
5use crate::core::pin::Pin;
6
7use crate::{Generator, GeneratorState};
8
9pub struct Iter<G>(G);
10
11impl<G> Iterator for Iter<G>
12where
13    G: Generator<(), Return = ()> + Unpin,
14{
15    type Item = G::Yield;
16
17    fn next(&mut self) -> Option<Self::Item> {
18        loop {
19            match Pin::new(&mut self.0).resume(()) {
20                GeneratorState::Complete(_) => return None,
21                GeneratorState::Suspend => continue,
22                GeneratorState::Yield(val) => return Some(val),
23            }
24        }
25    }
26}
27
28pub fn from_iter<Item>(
29    iter: impl IntoIterator<Item = Item>,
30) -> impl Generator<Yield = Item, Return = ()> {
31    struct GenIter<I: Iterator>(I);
32    impl<I: Iterator> Generator for GenIter<I> {
33        type Return = ();
34        type Yield = I::Item;
35        fn resume(self: Pin<&mut Self>, _value: ()) -> GeneratorState<Self::Yield, Self::Return> {
36            // SAFETY: We are not moving out of the pinned field.
37            match unsafe { self.get_unchecked_mut() }.0.next() {
38                Some(val) => GeneratorState::Yield(val),
39                None => GeneratorState::Complete(()),
40            }
41        }
42    }
43    GenIter(iter.into_iter())
44}
45
46pub const fn from_fn<Resume, Yield, Return>(
47    f: impl FnMut(Resume) -> GeneratorState<Yield, Return>,
48) -> impl Generator<Resume, Yield = Yield, Return = Return> {
49    struct GenFn<F>(F);
50    impl<F, R, Y, Out> Generator<R> for GenFn<F>
51    where
52        F: FnMut(R) -> GeneratorState<Y, Out>,
53    {
54        type Return = Out;
55        type Yield = Y;
56        fn resume(self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
57            // SAFETY: We are not moving out of the pinned field.
58            unsafe { self.get_unchecked_mut() }.0(value)
59        }
60    }
61    GenFn(f)
62}
63
64pub trait GeneratorIterator {
65    type Item;
66    type Iter: Iterator<Item = Self::Item>;
67    fn into_iter(self) -> Self::Iter;
68}
69
70impl<G> GeneratorIterator for G
71where
72    G: Generator<(), Return = ()> + Unpin,
73{
74    type Item = G::Yield;
75    type Iter = Iter<G>;
76    fn into_iter(self) -> Self::Iter {
77        Iter(self)
78    }
79}
80
81pub use either::Either;
82
83pub trait GeneratorExt<R>: Generator<R> {
84    fn receiving<F, U>(self, f: F) -> Receiving<Self, F>
85    where
86        Self: Sized,
87        F: FnMut(U) -> R,
88    {
89        Receiving { g: self, f }
90    }
91
92    fn compose<G: Generator<Self::Yield>>(self, then: G) -> Compose<Self, G>
93    where
94        Self: Sized,
95    {
96        Compose {
97            first: self,
98            second: then,
99        }
100    }
101
102    fn map_complete<U, F>(self, f: F) -> MapComplete<Self, F>
103    where
104        Self: Sized,
105        F: FnOnce(Self::Return) -> U,
106    {
107        MapComplete {
108            generator: self,
109            f: Some(f),
110        }
111    }
112
113    fn map_yield<U, F>(self, f: F) -> MapYield<Self, F>
114    where
115        Self: Sized,
116        F: FnMut(Self::Yield) -> U,
117    {
118        MapYield { generator: self, f }
119    }
120
121    fn and_then<G, F>(self, f: F) -> AndThen<Self, F, G>
122    where
123        Self: Sized,
124        G: Generator<R, Yield = Self::Yield>,
125        F: FnOnce(Self::Return) -> G,
126    {
127        AndThen::Before {
128            g: self,
129            f: Some(f),
130        }
131    }
132
133    fn flatten(self) -> Flatten<Self, Self::Yield>
134    where
135        Self: Sized,
136        Self::Yield: Generator<R, Return = ()>,
137    {
138        Flatten {
139            g: self,
140            current: None,
141        }
142    }
143}
144
145impl<G, R> GeneratorExt<R> for G where G: Generator<R> {}
146
147pub struct Compose<A, B> {
148    first: A,
149    second: B,
150}
151
152impl<A, B, R> Generator<R> for Compose<A, B>
153where
154    A: Generator<R>,
155    B: Generator<A::Yield>,
156{
157    type Return = Either<A::Return, B::Return>;
158    type Yield = B::Yield;
159
160    fn resume(mut self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
161        let first = unsafe { self.as_mut().map_unchecked_mut(|this| &mut this.first) };
162        let value = match first.resume(value) {
163            GeneratorState::Suspend => return GeneratorState::Suspend,
164            GeneratorState::Complete(x) => return GeneratorState::Complete(Either::Left(x)),
165            GeneratorState::Yield(value) => value,
166        };
167        let second = unsafe { self.map_unchecked_mut(|this| &mut this.second) };
168        second.resume(value).map_complete(Either::Right)
169    }
170}
171
172pub struct MapComplete<G, F> {
173    generator: G,
174    f: Option<F>,
175}
176
177pub struct MapYield<G, F> {
178    generator: G,
179    f: F,
180}
181
182impl<G, F, R, U> Generator<R> for MapComplete<G, F>
183where
184    G: Generator<R>,
185    F: FnOnce(G::Return) -> U,
186{
187    type Yield = G::Yield;
188    type Return = U;
189
190    fn resume(mut self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
191        let g = unsafe { self.as_mut().map_unchecked_mut(|this| &mut this.generator) };
192        g.resume(value)
193            .map_complete(|r| unsafe { self.get_unchecked_mut() }.f.take().unwrap()(r))
194    }
195}
196
197impl<G, F, R, U> Generator<R> for MapYield<G, F>
198where
199    G: Generator<R>,
200    F: FnMut(G::Yield) -> U,
201{
202    type Yield = U;
203    type Return = G::Return;
204
205    fn resume(self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
206        // safety: we are not moving out from the reference
207        let this = unsafe { self.get_unchecked_mut() };
208        let g = unsafe { Pin::new_unchecked(&mut this.generator) };
209        g.resume(value).map_yield(&mut this.f)
210    }
211}
212
213pub enum AndThen<G1, F, G2> {
214    Before { g: G1, f: Option<F> },
215    After { g: G2 },
216}
217
218impl<G1, F, G2, R> Generator<R> for AndThen<G1, F, G2>
219where
220    G1: Generator<R>,
221    F: FnOnce(G1::Return) -> G2,
222    G2: Generator<R, Yield = G1::Yield>,
223{
224    type Yield = G1::Yield;
225    type Return = G2::Return;
226
227    fn resume(mut self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
228        let this = unsafe { self.as_mut().get_unchecked_mut() };
229
230        match this {
231            AndThen::Before { g, f } => match unsafe { Pin::new_unchecked(g) }.resume(value) {
232                GeneratorState::Complete(r) => {
233                    let f = f.take().unwrap();
234                    let g = f(r);
235                    *this = AndThen::After { g };
236                    GeneratorState::Suspend
237                }
238                GeneratorState::Yield(y) => GeneratorState::Yield(y),
239                GeneratorState::Suspend => GeneratorState::Suspend,
240            },
241            AndThen::After { g } => unsafe { Pin::new_unchecked(g) }.resume(value),
242        }
243    }
244}
245
246pub struct Flatten<G, C> {
247    g: G,
248    current: Option<C>,
249}
250
251impl<R, G, C> Generator<R> for Flatten<G, C>
252where
253    G: Generator<R, Yield = C>,
254    C: Generator<R, Return = ()>,
255{
256    type Yield = C::Yield;
257    type Return = G::Return;
258
259    fn resume(self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
260        let this = unsafe { self.get_unchecked_mut() };
261
262        if let Some(current) = &mut this.current {
263            match unsafe { Pin::new_unchecked(current) }.resume(value) {
264                GeneratorState::Suspend => GeneratorState::Suspend,
265                GeneratorState::Complete(()) => {
266                    this.current = None;
267                    GeneratorState::Suspend
268                }
269                GeneratorState::Yield(y) => GeneratorState::Yield(y),
270            }
271        } else {
272            match unsafe { Pin::new_unchecked(&mut this.g) }.resume(value) {
273                GeneratorState::Complete(r) => GeneratorState::Complete(r),
274                GeneratorState::Yield(c) => {
275                    this.current = Some(c);
276                    GeneratorState::Suspend
277                }
278                GeneratorState::Suspend => GeneratorState::Suspend,
279            }
280        }
281    }
282}
283
284pub struct Once<Y>(Option<Y>);
285pub struct OnceWith<F>(Option<F>);
286
287pub struct CompleteWith<F, Y>(Option<F>, PhantomData<Y>);
288
289impl<Y> Generator for Once<Y> {
290    type Return = ();
291    type Yield = Y;
292    fn resume(self: Pin<&mut Self>, _value: ()) -> GeneratorState<Self::Yield, Self::Return> {
293        let this = unsafe { self.get_unchecked_mut() };
294        if let Some(y) = this.0.take() {
295            GeneratorState::Yield(y)
296        } else {
297            GeneratorState::Complete(())
298        }
299    }
300}
301
302impl<F, R, Y> Generator<R> for OnceWith<F>
303where
304    F: FnOnce(R) -> Y,
305{
306    type Return = ();
307    type Yield = Y;
308    fn resume(self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
309        let this = unsafe { self.get_unchecked_mut() };
310        if let Some(f) = this.0.take() {
311            GeneratorState::Yield(f(value))
312        } else {
313            GeneratorState::Complete(())
314        }
315    }
316}
317
318impl<F, R, C, Y> Generator<R> for CompleteWith<F, Y>
319where
320    F: FnOnce(R) -> C,
321{
322    type Return = C;
323    type Yield = Y;
324    fn resume(self: Pin<&mut Self>, value: R) -> GeneratorState<Self::Yield, Self::Return> {
325        let this = unsafe { self.get_unchecked_mut() };
326        let f = this.0.take().unwrap();
327        GeneratorState::Complete(f(value))
328    }
329}
330
331pub const fn identity<Y>() -> impl Generator<Y, Yield = Y, Return = Infallible> {
332    from_fn(GeneratorState::Yield)
333}
334
335pub const fn once<Y>(yielded: Y) -> impl Generator<Yield = Y, Return = ()> {
336    Once(Some(yielded))
337}
338pub const fn once_with<F, R, Y>(f: F) -> impl Generator<R, Yield = Y, Return = ()>
339where
340    F: FnOnce(R) -> Y,
341{
342    OnceWith(Some(f))
343}
344
345pub const fn complete_with<F, Y, R, C>(f: F) -> impl Generator<R, Yield = Y, Return = C>
346where
347    F: FnOnce(R) -> C,
348{
349    CompleteWith(Some(f), PhantomData)
350}
351
352pub struct Receiving<G, F> {
353    g: G,
354    f: F,
355}
356
357impl<R1, R2, G, F> Generator<R1> for Receiving<G, F>
358where
359    F: FnMut(R1) -> R2,
360    G: Generator<R2>,
361{
362    type Yield = G::Yield;
363    type Return = G::Return;
364
365    fn resume(self: Pin<&mut Self>, value: R1) -> GeneratorState<Self::Yield, Self::Return> {
366        let this = unsafe { self.get_unchecked_mut() };
367        let r2 = (this.f)(value);
368        unsafe { Pin::new_unchecked(&mut this.g) }.resume(r2)
369    }
370}
generator_light/ext.rs

generator_light/
ext.rs
