gat_lending_iterator/to_lending/
windows_mut.rs

1use crate::LendingIterator;
2
3/// A lending iterator over mutable windows.
4///
5/// This `struct` is created by the [`windows_mut`] method on [`ToLendingIterator`]. See
6/// its documentation for more.
7///
8/// [`ToLendingIterator`]: crate::ToLendingIterator
9/// [`windows_mut`]: crate::ToLendingIterator::windows_mut
10#[derive(Clone)]
11pub struct WindowsMut<I: Iterator> {
12    iter: I,
13    size: usize,
14    buf: Vec<I::Item>,
15}
16
17impl<I: Iterator> WindowsMut<I> {
18    pub(crate) fn new(mut iter: I, size: usize) -> Self {
19        let buf = iter.by_ref().take(size - 1).collect();
20        Self { iter, size, buf }
21    }
22}
23
24impl<I: Iterator> LendingIterator for WindowsMut<I> {
25    type Item<'a>
26        = &'a mut [I::Item]
27    where
28        Self: 'a;
29
30    fn next(&mut self) -> Option<Self::Item<'_>> {
31        self.iter.next().map(|next| {
32            if self.buf.len() == self.size * 2 - 1 {
33                self.buf.drain(..self.size);
34            }
35            self.buf.push(next);
36            let range = self.buf.len() - self.size..;
37            &mut self.buf[range]
38        })
39    }
40}
41
42#[cfg(test)]
43mod tests {
44    use crate::{LendingIterator, ToLendingIterator};
45
46    fn accumulate_window(w: &mut [i32]) -> i32 {
47        w[1] += w[0];
48        w[1]
49    }
50
51    #[test]
52    fn windows_mut_basic() {
53        let result: Vec<_> = (0..5)
54            .windows_mut(2)
55            .map(accumulate_window)
56            .into_iter()
57            .collect();
58        assert_eq!(result, vec![1, 3, 6, 10]);
59    }
60
61    #[test]
62    fn windows_mut_modifies_elements() {
63        let mut sum = 0;
64        (0..4).windows_mut(2).for_each(|w| {
65            w[0] *= 2;
66            sum += w[0];
67        });
68        assert_eq!(sum, 6);
69    }
70}