imbl_sized_chunks/sparse_chunk/
iter.rs

1use bitmaps::{Bitmap, Bits, BitsImpl};
2
3use super::SparseChunk;
4
5/// An iterator over references to the elements of a `SparseChunk`.
6pub struct Iter<'a, A, const N: usize>
7where
8    BitsImpl<N>: Bits,
9{
10    pub(crate) bitmap: Bitmap<N>,
11    pub(crate) chunk: &'a SparseChunk<A, N>,
12}
13
14// Implement Clone instead of deriving, because we want to be Clone even if A isn't.
15impl<A, const N: usize> Clone for Iter<'_, A, N>
16where
17    BitsImpl<N>: Bits,
18{
19    fn clone(&self) -> Self {
20        Iter {
21            bitmap: self.bitmap,
22            chunk: self.chunk,
23        }
24    }
25}
26
27impl<'a, A, const N: usize> Iterator for Iter<'a, A, N>
28where
29    BitsImpl<N>: Bits,
30{
31    type Item = &'a A;
32
33    fn next(&mut self) -> Option<Self::Item> {
34        if let Some(index) = self.bitmap.first_index() {
35            self.bitmap.set(index, false);
36            unsafe {
37                let p: *const A = self.chunk.get_unchecked(index);
38                Some(&*p)
39            }
40        } else {
41            None
42        }
43    }
44
45    fn size_hint(&self) -> (usize, Option<usize>) {
46        let len = self.bitmap.len();
47        (len, Some(len))
48    }
49}
50
51/// An iterator over mutable references to the elements of a `SparseChunk`.
52pub struct IterMut<'a, A, const N: usize>
53where
54    BitsImpl<N>: Bits,
55{
56    pub(crate) bitmap: Bitmap<N>,
57    pub(crate) chunk: &'a mut SparseChunk<A, N>,
58}
59
60impl<'a, A, const N: usize> Iterator for IterMut<'a, A, N>
61where
62    BitsImpl<N>: Bits,
63{
64    type Item = &'a mut A;
65
66    fn next(&mut self) -> Option<Self::Item> {
67        if let Some(index) = self.bitmap.first_index() {
68            self.bitmap.set(index, false);
69            unsafe {
70                let p: *mut A = self.chunk.get_unchecked_mut(index);
71                Some(&mut *p)
72            }
73        } else {
74            None
75        }
76    }
77
78    fn size_hint(&self) -> (usize, Option<usize>) {
79        let len = self.bitmap.len();
80        (len, Some(len))
81    }
82}
83
84/// A draining iterator over the elements of a `SparseChunk`.
85///
86/// "Draining" means that as the iterator yields each element, it's removed from
87/// the `SparseChunk`. When the iterator terminates, the chunk will be empty.
88pub struct Drain<A, const N: usize>
89where
90    BitsImpl<N>: Bits,
91{
92    pub(crate) chunk: SparseChunk<A, N>,
93}
94
95impl<A, const N: usize> Iterator for Drain<A, N>
96where
97    BitsImpl<N>: Bits,
98{
99    type Item = A;
100
101    fn next(&mut self) -> Option<Self::Item> {
102        self.chunk.pop()
103    }
104
105    fn size_hint(&self) -> (usize, Option<usize>) {
106        let len = self.chunk.len();
107        (len, Some(len))
108    }
109}
110
111/// An iterator over `Option`s of references to the elements of a `SparseChunk`.
112///
113/// Iterates over every index in the `SparseChunk`, from zero to its full capacity,
114/// returning an `Option<&A>` for each index.
115pub struct OptionIter<'a, A, const N: usize>
116where
117    BitsImpl<N>: Bits,
118{
119    pub(crate) index: usize,
120    pub(crate) chunk: &'a SparseChunk<A, N>,
121}
122
123impl<'a, A, const N: usize> Iterator for OptionIter<'a, A, N>
124where
125    BitsImpl<N>: Bits,
126{
127    type Item = Option<&'a A>;
128
129    fn next(&mut self) -> Option<Self::Item> {
130        if self.index < N {
131            let result = self.chunk.get(self.index);
132            self.index += 1;
133            Some(result)
134        } else {
135            None
136        }
137    }
138
139    fn size_hint(&self) -> (usize, Option<usize>) {
140        (
141            SparseChunk::<A, N>::CAPACITY - self.index,
142            Some(SparseChunk::<A, N>::CAPACITY - self.index),
143        )
144    }
145}
146
147/// An iterator over `Option`s of mutable references to the elements of a `SparseChunk`.
148///
149/// Iterates over every index in the `SparseChunk`, from zero to its full capacity,
150/// returning an `Option<&mut A>` for each index.
151pub struct OptionIterMut<'a, A, const N: usize>
152where
153    BitsImpl<N>: Bits,
154{
155    pub(crate) index: usize,
156    pub(crate) chunk: &'a mut SparseChunk<A, N>,
157}
158
159impl<'a, A, const N: usize> Iterator for OptionIterMut<'a, A, N>
160where
161    BitsImpl<N>: Bits,
162{
163    type Item = Option<&'a mut A>;
164
165    fn next(&mut self) -> Option<Self::Item> {
166        if self.index < N {
167            let result = if self.chunk.map.get(self.index) {
168                unsafe {
169                    let p: *mut A = &mut self.chunk.values_mut()[self.index];
170                    Some(Some(&mut *p))
171                }
172            } else {
173                Some(None)
174            };
175            self.index += 1;
176            result
177        } else {
178            None
179        }
180    }
181
182    fn size_hint(&self) -> (usize, Option<usize>) {
183        (
184            SparseChunk::<A, N>::CAPACITY - self.index,
185            Some(SparseChunk::<A, N>::CAPACITY - self.index),
186        )
187    }
188}
189
190/// A draining iterator over `Option`s of the elements of a `SparseChunk`.
191///
192/// Iterates over every index in the `SparseChunk`, from zero to its full capacity,
193/// returning an `Option<A>` for each index.
194pub struct OptionDrain<A, const N: usize>
195where
196    BitsImpl<N>: Bits,
197{
198    pub(crate) index: usize,
199    pub(crate) chunk: SparseChunk<A, N>,
200}
201
202impl<A, const N: usize> Iterator for OptionDrain<A, N>
203where
204    BitsImpl<N>: Bits,
205{
206    type Item = Option<A>;
207
208    fn next(&mut self) -> Option<Self::Item> {
209        if self.index < N {
210            let result = self.chunk.remove(self.index);
211            self.index += 1;
212            Some(result)
213        } else {
214            None
215        }
216    }
217
218    fn size_hint(&self) -> (usize, Option<usize>) {
219        (
220            SparseChunk::<A, N>::CAPACITY - self.index,
221            Some(SparseChunk::<A, N>::CAPACITY - self.index),
222        )
223    }
224}
225
226#[cfg(test)]
227mod test {
228    use super::*;
229    use std::iter::FromIterator;
230
231    #[test]
232    fn iter() {
233        let vec: Vec<Option<usize>> =
234            Vec::from_iter((0..64).map(|i| if i % 2 == 0 { Some(i) } else { None }));
235        let chunk: SparseChunk<usize, 64> = vec.iter().cloned().collect();
236        let vec: Vec<usize> = vec.iter().cloned().flatten().collect();
237        assert!(vec.iter().eq(chunk.iter()));
238    }
239
240    #[test]
241    fn iter_mut() {
242        let vec: Vec<Option<usize>> =
243            Vec::from_iter((0..64).map(|i| if i % 2 == 0 { Some(i) } else { None }));
244        let mut chunk: SparseChunk<_, 64> = vec.iter().cloned().collect();
245        let mut vec: Vec<usize> = vec.iter().cloned().flatten().collect();
246        assert!(vec.iter_mut().eq(chunk.iter_mut()));
247    }
248
249    #[test]
250    fn drain() {
251        let vec: Vec<Option<usize>> =
252            Vec::from_iter((0..64).map(|i| if i % 2 == 0 { Some(i) } else { None }));
253        let chunk: SparseChunk<_, 64> = vec.iter().cloned().collect();
254        let vec: Vec<usize> = vec.iter().cloned().flatten().collect();
255        assert!(vec.into_iter().eq(chunk.into_iter()));
256    }
257
258    #[test]
259    fn option_iter() {
260        let vec: Vec<Option<usize>> =
261            Vec::from_iter((0..64).map(|i| if i % 2 == 0 { Some(i) } else { None }));
262        let chunk: SparseChunk<_, 64> = vec.iter().cloned().collect();
263        assert!(vec
264            .iter()
265            .cloned()
266            .eq(chunk.option_iter().map(|v| v.cloned())));
267    }
268
269    #[test]
270    fn option_iter_mut() {
271        let vec: Vec<Option<usize>> =
272            Vec::from_iter((0..64).map(|i| if i % 2 == 0 { Some(i) } else { None }));
273        let mut chunk: SparseChunk<_, 64> = vec.iter().cloned().collect();
274        assert!(vec
275            .iter()
276            .cloned()
277            .eq(chunk.option_iter_mut().map(|v| v.cloned())));
278    }
279
280    #[test]
281    fn option_drain() {
282        let vec: Vec<Option<usize>> =
283            Vec::from_iter((0..64).map(|i| if i % 2 == 0 { Some(i) } else { None }));
284        let chunk: SparseChunk<_, 64> = vec.iter().cloned().collect();
285        assert!(vec.iter().cloned().eq(chunk.option_drain()));
286    }
287}