iterators_extended/
stateful.rs

1pub struct StatefulMapIterator<I, S, F> {
2    state: S,
3    iter: I,
4    func: F,
5}
6
7impl<I: Iterator, S, F, FO> Iterator for StatefulMapIterator<I, S, F>
8where
9    F: FnMut(&mut S, I::Item) -> FO,
10{
11    type Item = FO;
12
13    fn next(&mut self) -> Option<Self::Item> {
14        self.iter.next().map(|x| (self.func)(&mut self.state, x))
15    }
16}
17
18pub trait IntoStatefulMapIterator: IntoIterator {
19    fn stateful_map<S, F>(self, func: F, state: S) -> StatefulMapIterator<Self::IntoIter, S, F>
20    where
21        Self: Sized,
22    {
23        StatefulMapIterator {
24            iter: self.into_iter(),
25            state,
26            func,
27        }
28    }
29}
30
31impl<I: IntoIterator> IntoStatefulMapIterator for I {}
32
33#[cfg(test)]
34mod tests {
35    use crate::stateful::IntoStatefulMapIterator;
36
37    #[test]
38    fn test_simple_stateful() {
39        struct State {
40            total: u8,
41        }
42        fn cumsum(state: &mut State, x: u8) -> u8 {
43            state.total += x;
44            state.total
45        }
46
47        let values = (0_u8..16).collect::<Vec<_>>();
48        let mapped = values
49            .stateful_map(cumsum, State { total: 0 })
50            .collect::<Vec<_>>();
51
52        println!("{:?}", mapped);
53    }
54
55    #[test]
56    fn test_fib() {
57        struct State {
58            prev0: usize,
59            prev1: usize,
60        }
61        fn fib(state: &mut State, _x: usize) -> usize {
62            let next = state.prev0 + state.prev1;
63            state.prev0 = state.prev1;
64            state.prev1 = next;
65            next
66        }
67
68        let values = (0..16).collect::<Vec<_>>();
69        let mapped = values
70            .stateful_map(fib, State { prev0: 1, prev1: 0 })
71            .collect::<Vec<_>>();
72        println!("{:?}", mapped);
73    }
74
75    #[test]
76    fn test_reuse_alloc() {
77        struct State {
78            buffer: Vec<u8>,
79        }
80        fn arbitrary_vector_stuff(state: &mut State, _x: u8) -> u8 {
81            // Simulate some expensive operations on an array
82            for item in state.buffer.iter_mut() {
83                for _ in 0..100 {
84                    *item += 1;
85                    *item -= 1;
86                }
87            }
88            _x
89        }
90
91        let values = (0..16).collect::<Vec<_>>();
92        let mapped = values
93            .stateful_map(
94                arbitrary_vector_stuff,
95                State {
96                    buffer: vec![0; 1000000],
97                },
98            )
99            .collect::<Vec<_>>();
100        println!("{:?}", mapped);
101    }
102}
iterators_extended/stateful.rs

iterators_extended/
stateful.rs
