rollercoaster/
memory.rs

1pub struct Memory<I>
2where
3    I: Iterator,
4{
5    items: Vec<I::Item>,
6    underlying: I,
7}
8
9impl<I> Memory<I>
10where
11    I: Iterator,
12{
13    /// Remember this item for the next iteration.
14    pub fn remember(&mut self, item: I::Item) {
15        self.items.push(item);
16    }
17
18    /// Clears all items that were remembered,
19    /// and the iterator will now return items
20    /// from the underlying iterator.
21    pub fn clear(&mut self) {
22        self.items.clear();
23    }
24
25    /// Removes the most recently remembered item.
26    pub fn forget(&mut self) {
27        self.items.pop();
28    }
29
30    pub(crate) fn new(underlying: I) -> Self {
31        Self {
32            underlying,
33            items: vec![],
34        }
35    }
36}
37
38impl<I> Iterator for Memory<I>
39where
40    I: Iterator,
41{
42    type Item = I::Item;
43
44    fn next(&mut self) -> Option<Self::Item> {
45        if !self.items.is_empty() {
46            return self.items.pop();
47        }
48
49        self.underlying.next()
50    }
51}
52
53ext_impl! {
54    /// Creates an iterator that allows remembering one or more values
55    /// for the next iterations.
56    ///
57    /// ## Why?
58    /// Sometimes you are working with an iterator in a way
59    /// where you need to own values, and at the same time
60    /// conditionally perform some action on them across iterations.
61    ///
62    /// As an example, this is used in
63    /// [`group_by()`](crate::Rollercoaster::group_by) to allow lazy grouping.
64    ///
65    /// ## How is this different from `peekable()`?
66    /// Unlike [`Peekable`](std::iter::Peekable),
67    /// it allows you to own the value from the current iteration
68    /// and choose to use it again in the next one.
69    ///
70    /// Memory also works in the opposite way, there is no peeking
71    /// and instead you can check some condition in the _next_ iteration.
72    ///
73    /// # Example
74    /// ```
75    /// # use rollercoaster::Rollercoaster;
76    /// #
77    /// let mut nums = vec![1, 2, 3, 4, 5, 3].into_iter().memory();
78    ///
79    /// for n in nums.by_ref() {
80    ///     // When 4 is encountered, we want to sum it with
81    ///     // everything after it.
82    ///     if n == 4 {
83    ///         nums.remember(n);
84    ///         break;
85    ///     }
86    /// }
87    ///
88    /// let summed: u32 = nums.sum();
89    /// assert_eq!(summed, 12);
90    /// ```
91    fn memory(self) -> Memory<Self> {
92        Memory::new(self)
93    }
94}
95
96#[cfg(test)]
97mod test {
98    use crate::Rollercoaster;
99
100    fn mock() -> Vec<&'static str> {
101        vec!["a", "b", "c", "d", "e"]
102    }
103
104    #[test]
105    fn it_remembers() {
106        let mut memory = mock().into_iter().memory();
107
108        for letter in memory.by_ref() {
109            if letter == "d" {
110                memory.remember(letter);
111                break;
112            }
113        }
114
115        assert_eq!(memory.collect::<Vec<_>>(), vec!["d", "e"]);
116    }
117
118    #[test]
119    fn it_forgets() {
120        let mut memory = mock().into_iter().memory();
121
122        memory.remember("f");
123        memory.remember("g");
124        memory.forget();
125
126        assert_eq!(
127            memory.collect::<Vec<_>>(),
128            vec!["f", "a", "b", "c", "d", "e"]
129        );
130    }
131}