genawaiter/stack/
iterator.rs

1use crate::{ops::GeneratorState, stack::generator::Gen};
2use std::future::Future;
3
4impl<'s, Y, F: Future<Output = ()>> IntoIterator for Gen<'s, Y, (), F> {
5    type Item = Y;
6    type IntoIter = IntoIter<'s, Y, F>;
7
8    #[must_use]
9    fn into_iter(self) -> Self::IntoIter {
10        IntoIter { generator: self }
11    }
12}
13
14pub struct IntoIter<'s, Y, F: Future<Output = ()>> {
15    generator: Gen<'s, Y, (), F>,
16}
17
18impl<'s, Y, F: Future<Output = ()>> Iterator for IntoIter<'s, Y, F> {
19    type Item = Y;
20
21    fn next(&mut self) -> Option<Self::Item> {
22        match self.generator.resume() {
23            GeneratorState::Yielded(x) => Some(x),
24            GeneratorState::Complete(()) => None,
25        }
26    }
27}
28
29impl<'r, 's, Y, F: Future<Output = ()>> IntoIterator for &'r mut Gen<'s, Y, (), F> {
30    type Item = Y;
31    type IntoIter = MutIntoIter<'r, 's, Y, F>;
32
33    fn into_iter(self) -> Self::IntoIter {
34        MutIntoIter { generator: self }
35    }
36}
37
38pub struct MutIntoIter<'r, 's, Y, F: Future<Output = ()>> {
39    generator: &'r mut Gen<'s, Y, (), F>,
40}
41
42impl<'r, 's, Y, F: Future<Output = ()>> Iterator for MutIntoIter<'r, 's, Y, F> {
43    type Item = Y;
44
45    fn next(&mut self) -> Option<Self::Item> {
46        match self.generator.resume() {
47            GeneratorState::Yielded(x) => Some(x),
48            GeneratorState::Complete(()) => None,
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use crate::stack::{let_gen_using, Co, Gen, Shelf};
56    use std::iter::IntoIterator;
57
58    async fn produce(co: Co<'_, i32>) {
59        co.yield_(10).await;
60        co.yield_(20).await;
61    }
62
63    #[test]
64    fn let_gen_using_into_iter() {
65        let_gen_using!(gen, produce);
66
67        let items: Vec<_> = gen.into_iter().collect();
68        assert_eq!(items, [10, 20]);
69    }
70
71    #[test]
72    fn let_gen_using_for_loop() {
73        let_gen_using!(gen, produce);
74
75        let mut sum = 0;
76        for x in gen {
77            sum += x;
78        }
79        assert_eq!(sum, 30);
80    }
81
82    #[test]
83    fn shelf_generator_into_iter() {
84        let mut shelf = Shelf::new();
85        let gen = unsafe { Gen::new(&mut shelf, produce) };
86
87        let items: Vec<_> = gen.into_iter().collect();
88        assert_eq!(items, [10, 20]);
89    }
90
91    #[test]
92    fn shelf_generator_for_loop() {
93        let mut shelf = Shelf::new();
94        let gen = unsafe { Gen::new(&mut shelf, produce) };
95
96        let mut sum = 0;
97        for x in gen {
98            sum += x;
99        }
100        assert_eq!(sum, 30);
101    }
102}
genawaiter/stack/iterator.rs

genawaiter/stack/
iterator.rs
