generate/
iterator.rs

1use std::pin::Pin;
2
3use crate::{Generator, GeneratorState};
4
5struct Iter<G> {
6    generator: Pin<Box<G>>,
7}
8
9/// Turns the given generator into an iterator.
10///
11/// The returned iterator iterates over `Generator::Yield` and returns `None` when
12/// `GeneratorState::Complete(_)` is returned.
13pub fn iterate<G>(generator: G) -> impl Iterator<Item = G::Yield>
14where
15    G: Generator,
16{
17    Iter {
18        generator: Box::pin(generator),
19    }
20}
21
22impl<G> Iterator for Iter<G>
23where
24    G: Generator,
25{
26    type Item = G::Yield;
27
28    fn next(&mut self) -> Option<Self::Item> {
29        match self.generator.as_mut().resume(()) {
30            GeneratorState::Yielded(val) => Some(val),
31            GeneratorState::Complete(_) => None,
32        }
33    }
34}
35
36struct IterWithReturn<G, F> {
37    generator: Pin<Box<G>>,
38    handle_return: Option<F>,
39}
40/// Turns the given generator into an iterator, with a function to handle the
41/// `GeneratorState::Complete(_)` case.
42///
43/// The returned iterator iterates over `Generator::Yield`. When `GeneratorState::Complete(_)`
44/// is returned, it calls `handle_return` with the returned value. The return of that call is then
45/// emitted from the iterator. All subsequent calls to `Iterator::next` then return `None`.
46pub fn iterate_with_return<G, F>(generator: G, handle_return: F) -> impl Iterator<Item = G::Yield>
47where
48    G: Generator,
49    F: FnOnce(G::Return) -> Option<G::Yield>,
50{
51    IterWithReturn {
52        generator: Box::pin(generator),
53        handle_return: Some(handle_return),
54    }
55}
56
57impl<G, F> Iterator for IterWithReturn<G, F>
58where
59    G: Generator,
60    F: FnOnce(G::Return) -> Option<G::Yield>,
61{
62    type Item = G::Yield;
63
64    fn next(&mut self) -> Option<Self::Item> {
65        if self.handle_return.is_none() {
66            return None;
67        }
68        match self.generator.as_mut().resume(()) {
69            GeneratorState::Yielded(val) => Some(val),
70            GeneratorState::Complete(val) => {
71                if let Some(handle_return) = self.handle_return.take() {
72                    handle_return(val)
73                } else {
74                    None
75                }
76            }
77        }
78    }
79}
generate/iterator.rs

generate/
iterator.rs
