streaming_iterator/
slice.rs

1use crate::{
2    DoubleEndedStreamingIterator, DoubleEndedStreamingIteratorMut, StreamingIterator,
3    StreamingIteratorMut,
4};
5
6use core::mem;
7use core::num::NonZeroUsize;
8
9/// Creates an iterator over all contiguous windows of length `size` in a mutable `slice`.
10///
11/// The windows overlap and may be mutated via `StreamingIteratorMut`.
12/// If the `slice` is shorter than `size`, the iterator returns no values.
13///
14/// # Panics
15///
16/// Panics if `size` is 0.
17pub fn windows_mut<T>(slice: &mut [T], size: usize) -> WindowsMut<'_, T> {
18    WindowsMut {
19        slice,
20        size: NonZeroUsize::new(size).expect("size is zero"),
21        position: Position::Init,
22    }
23}
24
25/// A streaming iterator which returns overlapping mutable subslices of length `size`.
26///
27/// This struct is created by the [`windows_mut`] function.
28pub struct WindowsMut<'a, T> {
29    slice: &'a mut [T],
30    size: NonZeroUsize,
31    position: Position,
32}
33
34enum Position {
35    Init,
36    Front,
37    Back,
38}
39
40impl<T> WindowsMut<'_, T> {
41    fn consume(&mut self) {
42        match self.position {
43            Position::Init => {}
44            Position::Front => {
45                let slice = mem::take(&mut self.slice);
46                if let Some((_, tail)) = slice.split_first_mut() {
47                    self.slice = tail;
48                }
49            }
50            Position::Back => {
51                let slice = mem::take(&mut self.slice);
52                if let Some((_, head)) = slice.split_last_mut() {
53                    self.slice = head;
54                }
55            }
56        }
57    }
58
59    fn get_front(&self) -> Option<&[T]> {
60        self.slice.get(..self.size.get())
61    }
62
63    fn get_front_mut(&mut self) -> Option<&mut [T]> {
64        self.slice.get_mut(..self.size.get())
65    }
66
67    fn get_back(&self) -> Option<&[T]> {
68        let start = self.slice.len().checked_sub(self.size.get())?;
69        self.slice.get(start..)
70    }
71
72    fn get_back_mut(&mut self) -> Option<&mut [T]> {
73        let start = self.slice.len().checked_sub(self.size.get())?;
74        self.slice.get_mut(start..)
75    }
76
77    fn len(&self) -> usize {
78        let len = match self.position {
79            Position::Init => self.slice.len(),
80            _ => self.slice.len().saturating_sub(1),
81        };
82        len.saturating_sub(self.size.get() - 1)
83    }
84}
85
86impl<T> StreamingIterator for WindowsMut<'_, T> {
87    type Item = [T];
88
89    fn advance(&mut self) {
90        self.consume();
91        self.position = Position::Front;
92    }
93
94    fn get(&self) -> Option<&Self::Item> {
95        match self.position {
96            Position::Init => None,
97            Position::Front => self.get_front(),
98            Position::Back => self.get_back(),
99        }
100    }
101
102    fn next(&mut self) -> Option<&Self::Item> {
103        self.advance();
104        self.get_front()
105    }
106
107    fn size_hint(&self) -> (usize, Option<usize>) {
108        let len = self.len();
109        (len, Some(len))
110    }
111
112    fn is_done(&self) -> bool {
113        self.slice.len() < self.size.get()
114    }
115
116    fn count(self) -> usize {
117        self.len()
118    }
119}
120
121impl<T> StreamingIteratorMut for WindowsMut<'_, T> {
122    fn get_mut(&mut self) -> Option<&mut Self::Item> {
123        match self.position {
124            Position::Init => None,
125            Position::Front => self.get_front_mut(),
126            Position::Back => self.get_back_mut(),
127        }
128    }
129
130    fn next_mut(&mut self) -> Option<&mut Self::Item> {
131        self.advance();
132        self.get_front_mut()
133    }
134}
135
136impl<T> DoubleEndedStreamingIterator for WindowsMut<'_, T> {
137    fn advance_back(&mut self) {
138        self.consume();
139        self.position = Position::Back;
140    }
141
142    fn next_back(&mut self) -> Option<&Self::Item> {
143        self.advance_back();
144        self.get_back()
145    }
146}
147
148impl<T> DoubleEndedStreamingIteratorMut for WindowsMut<'_, T> {
149    fn next_back_mut(&mut self) -> Option<&mut Self::Item> {
150        self.advance_back();
151        self.get_back_mut()
152    }
153}
154
155#[test]
156fn test_windows_mut() {
157    let slice: &mut [_] = &mut [0; 6];
158
159    windows_mut(slice, 3).fold_mut(0, |i, win| {
160        win.copy_from_slice(&[i; 3]);
161        i + 1
162    });
163    assert_eq!(slice, &[0, 1, 2, 3, 3, 3]);
164
165    windows_mut(slice, 2).rfold_mut(0, |i, win| {
166        win.copy_from_slice(&[i; 2]);
167        i + 1
168    });
169    assert_eq!(slice, &[4, 4, 3, 2, 1, 0]);
170
171    let mut i = 0;
172    let mut iter = windows_mut(slice, 1);
173    while let Some(win) = iter.next_mut() {
174        win.copy_from_slice(&[i]);
175        i += 1;
176
177        if let Some(win) = iter.next_back_mut() {
178            win.copy_from_slice(&[i]);
179            i += 1;
180        }
181    }
182    assert_eq!(slice, &[0, 2, 4, 5, 3, 1]);
183}
184
185#[test]
186fn test_windows_mut_count() {
187    let slice: &mut [_] = &mut [0; 6];
188
189    assert_eq!(windows_mut(slice, 3).count(), 4);
190    assert_eq!(windows_mut(slice, 6).count(), 1);
191    assert_eq!(windows_mut(slice, 9).count(), 0);
192
193    let mut iter = windows_mut(slice, 3);
194    assert_eq!(iter.size_hint(), (4, Some(4)));
195    iter.advance();
196    assert_eq!(iter.count(), 3);
197}
198
199#[test]
200#[should_panic]
201fn test_windows_mut_0() {
202    let _: WindowsMut<'_, i32> = windows_mut(&mut [], 0);
203}
streaming_iterator/slice.rs

streaming_iterator/
slice.rs
