read_write_store/
iter.rs

1//! Iterator implementations for [RwStore](crate::RwStore).
2
3use std::panic::{RefUnwindSafe, UnwindSafe};
4use std::pin::Pin;
5
6use crate::block::Block;
7use crate::bucket::head_size_to_total_size;
8use crate::id::Id;
9use crate::util::debug_check::UnwrapDebugChecked;
10use crate::RwStore;
11use std::iter::FusedIterator;
12
13/// An owned iterator over an [RwStore](crate::RwStore).
14///
15/// # Example
16///
17/// ```
18/// # use read_write_store::RwStore;
19/// let store = RwStore::new();
20/// let id = store.insert(42);
21/// let mut iter = store.into_iter();
22/// assert_eq!(iter.next().unwrap().1, 42);
23/// assert!(iter.next().is_none());
24/// ```
25pub struct IntoIter<Element> {
26    iter: GenIter<OwnedBlockTracker<Element>>,
27}
28
29impl<Element> IntoIter<Element> {
30    /// Creates an owned iterator over an [RwStore](crate::RwStore).
31    pub fn new(store: RwStore<Element>) -> Self {
32        let bucket_iters = IntoIterator::into_iter(store.buckets)
33            .map(|bucket| {
34                let block_list = bucket.into_inner().inner.into_inner();
35
36                let last_slot = block_list.next_unused_slot.into_inner();
37                let current_block = block_list.head.map(OwnedBlockTracker::new);
38
39                GenBucketIter::new(last_slot, current_block)
40            })
41            .collect();
42
43        Self {
44            iter: GenIter::new(bucket_iters),
45        }
46    }
47}
48
49impl<Element> Iterator for IntoIter<Element> {
50    type Item = (Id, Element);
51
52    fn next(&mut self) -> Option<(Id, Element)> {
53        self.iter.next()
54    }
55
56    fn size_hint(&self) -> (usize, Option<usize>) {
57        self.iter.size_hint()
58    }
59}
60
61impl<Element> FusedIterator for IntoIter<Element> {}
62
63unsafe impl<Element> Sync for IntoIter<Element> {}
64
65unsafe impl<Element: Send> Send for IntoIter<Element> {}
66
67impl<Element: UnwindSafe> UnwindSafe for IntoIter<Element> {}
68
69impl<Element: RefUnwindSafe> RefUnwindSafe for IntoIter<Element> {}
70
71/// A mutable iterator over an [RwStore](crate::RwStore).
72///
73/// # Example
74///
75/// ```
76/// # use read_write_store::RwStore;
77/// let mut store = RwStore::new();
78/// let id = store.insert(42);
79/// let mut iter = store.iter_mut();
80/// assert_eq!(iter.next().unwrap().1, &mut 42);
81/// assert!(iter.next().is_none());
82/// ```
83pub struct IterMut<'a, Element> {
84    iter: GenIter<MutBlockTracker<'a, Element>>,
85}
86
87impl<'a, Element> IterMut<'a, Element> {
88    /// Creates a mutable iterator over an [RwStore](crate::RwStore).
89    pub fn new(store: &'a mut RwStore<Element>) -> Self {
90        let bucket_iters = store
91            .buckets
92            .iter_mut()
93            .map(|bucket| {
94                let block_list = bucket.inner.get_mut();
95
96                let last_slot = block_list.next_unused_slot.load_directly();
97                let current_block = block_list.head.as_mut().map(MutBlockTracker::new);
98
99                GenBucketIter::new(last_slot, current_block)
100            })
101            .collect();
102
103        Self {
104            iter: GenIter::new(bucket_iters),
105        }
106    }
107}
108
109impl<'a, Element> Iterator for IterMut<'a, Element> {
110    type Item = (Id, &'a mut Element);
111
112    fn next(&mut self) -> Option<(Id, &'a mut Element)> {
113        self.iter.next()
114    }
115
116    fn size_hint(&self) -> (usize, Option<usize>) {
117        self.iter.size_hint()
118    }
119}
120
121impl<'a, Element> FusedIterator for IterMut<'a, Element> {}
122
123unsafe impl<'a, Element> Sync for IterMut<'a, Element> {}
124
125unsafe impl<'a, Element: Send> Send for IterMut<'a, Element> {}
126
127impl<'a, Element: UnwindSafe> UnwindSafe for IterMut<'a, Element> {}
128
129impl<'a, Element: RefUnwindSafe> RefUnwindSafe for IterMut<'a, Element> {}
130
131struct GenIter<Tracker: BlockTracker> {
132    iters: Vec<GenBucketIter<Tracker>>,
133}
134
135impl<Tracker: BlockTracker> GenIter<Tracker> {
136    pub fn new(iters: Vec<GenBucketIter<Tracker>>) -> Self {
137        Self { iters }
138    }
139}
140
141impl<Tracker: BlockTracker> Iterator for GenIter<Tracker> {
142    type Item = Tracker::Element;
143
144    fn next(&mut self) -> Option<Self::Item> {
145        let current = self.iters.last_mut()?;
146
147        if let Some(result) = current.next() {
148            Some(result)
149        } else {
150            self.iters.pop();
151            self.next()
152        }
153    }
154
155    fn size_hint(&self) -> (usize, Option<usize>) {
156        let (mut total_lower, mut total_upper) = (0, Some(0));
157
158        for iter in &self.iters {
159            let (lower, upper) = iter.size_hint();
160
161            total_lower += lower;
162            total_upper = total_upper.zip_with(upper, |a, b| a + b);
163        }
164
165        (total_lower, total_upper)
166    }
167}
168
169impl<Tracker: BlockTracker> FusedIterator for GenIter<Tracker> {}
170
171struct GenBucketIter<Tracker: BlockTracker> {
172    last_slot: u32,
173    current_block: Option<Tracker>,
174}
175
176impl<Tracker: BlockTracker> GenBucketIter<Tracker> {
177    pub fn new(last_slot: u32, current_block: Option<Tracker>) -> Self {
178        let max_last_slot = current_block.as_ref().map(|block| block.len()).unwrap_or(0);
179        let last_slot = last_slot.min(max_last_slot);
180
181        Self {
182            last_slot,
183            current_block,
184        }
185    }
186}
187
188impl<Tracker: BlockTracker> Iterator for GenBucketIter<Tracker> {
189    type Item = Tracker::Element;
190
191    fn next(&mut self) -> Option<Tracker::Element> {
192        unsafe {
193            if self.last_slot == 0 {
194                let old_block = self.current_block.take()?;
195
196                let new_block = old_block.next();
197                let new_block_size = new_block.as_ref().map(|block| block.len()).unwrap_or(0);
198
199                self.current_block = new_block;
200                self.last_slot = new_block_size;
201
202                return self.next();
203            }
204
205            let current_block = self.current_block.as_mut().unwrap_debug_checked();
206            self.last_slot -= 1;
207
208            let contents = current_block.take(self.last_slot);
209            match contents {
210                Some(result) => Some(result),
211                None => self.next(),
212            }
213        }
214    }
215
216    fn size_hint(&self) -> (usize, Option<usize>) {
217        let current_block_size = self
218            .current_block
219            .as_ref()
220            .map(|block| block.len())
221            .unwrap_or(0);
222
223        let allocated_size = head_size_to_total_size(current_block_size);
224        let touched_size = allocated_size - (current_block_size - self.last_slot);
225
226        (0, Some(touched_size as usize))
227    }
228}
229
230impl<Tracker: BlockTracker> FusedIterator for GenBucketIter<Tracker> {}
231
232trait BlockTracker: Sized {
233    type Element;
234
235    fn next(self) -> Option<Self>;
236
237    unsafe fn take(&mut self, slot: u32) -> Option<Self::Element>;
238
239    fn len(&self) -> u32;
240}
241
242struct OwnedBlockTracker<Element> {
243    block: Pin<Box<Block<Element>>>,
244}
245
246impl<Element> OwnedBlockTracker<Element> {
247    pub fn new(block: Pin<Box<Block<Element>>>) -> Self {
248        Self { block }
249    }
250}
251
252impl<Element> BlockTracker for OwnedBlockTracker<Element> {
253    type Element = (Id, Element);
254
255    fn next(self) -> Option<Self> {
256        Block::into_next(self.block).map(OwnedBlockTracker::new)
257    }
258
259    unsafe fn take(&mut self, slot: u32) -> Option<(Id, Element)> {
260        self.block
261            .as_mut()
262            .get_unchecked_mut()
263            .take(slot)
264    }
265
266    fn len(&self) -> u32 {
267        self.block.len()
268    }
269}
270
271struct MutBlockTracker<'a, Element> {
272    block: &'a mut Pin<Box<Block<Element>>>,
273}
274
275impl<'a, Element> MutBlockTracker<'a, Element> {
276    pub fn new(block: &'a mut Pin<Box<Block<Element>>>) -> Self {
277        Self { block }
278    }
279}
280
281impl<'a, Element> BlockTracker for MutBlockTracker<'a, Element> {
282    type Element = (Id, &'a mut Element);
283
284    fn next(self) -> Option<Self> {
285        unsafe {
286            self.block
287                .as_mut()
288                .get_unchecked_mut()
289                .next_mut()
290                .map(MutBlockTracker::new)
291        }
292    }
293
294    unsafe fn take(&mut self, slot: u32) -> Option<(Id, &'a mut Element)> {
295        self.block
296            .as_mut()
297            .get_unchecked_mut()
298            .slot_contents(slot)
299    }
300
301    fn len(&self) -> u32 {
302        self.block.len()
303    }
304}
305
306#[cfg(test)]
307mod test {
308    use std::panic::{RefUnwindSafe, UnwindSafe};
309
310    use crate::RwStore;
311
312    #[test]
313    fn into_iter_has_correct_size() {
314        iter_has_correct_size(|store| store.into_iter().count());
315    }
316
317    #[test]
318    fn iter_mut_has_correct_size() {
319        iter_has_correct_size(|mut store| store.iter_mut().count());
320    }
321
322    fn iter_has_correct_size(count: impl Fn(RwStore<u32>) -> usize) {
323        let store = RwStore::new();
324        for i in 0..42 {
325            store.insert(i);
326        }
327        assert_eq!(count(store), 42);
328
329        let store = RwStore::new();
330        store.insert(42);
331        assert_eq!(count(store), 1);
332
333        let store = RwStore::new();
334        assert_eq!(count(store), 0);
335    }
336
337    #[test]
338    fn into_iter_has_correct_elements() {
339        iter_has_correct_elements(|store| store.into_iter().map(|(_, value)| value).sum())
340    }
341
342    #[test]
343    fn iter_mut_has_correct_elements() {
344        iter_has_correct_elements(|mut store| store.iter_mut().map(|(_, value)| *value).sum())
345    }
346
347    fn iter_has_correct_elements(sum: impl Fn(RwStore<u32>) -> u32) {
348        let store = RwStore::new();
349        store.insert(1);
350        store.insert(2);
351        store.insert(4);
352        assert_eq!(sum(store), 7);
353
354        let store = RwStore::new();
355        for i in 1..=42 {
356            store.insert(i);
357        }
358        assert_eq!(sum(store), 903);
359
360        let store = RwStore::new();
361        assert_eq!(sum(store), 0);
362    }
363
364    #[test]
365    fn into_iter_skips_removed_elements() {
366        iter_skips_removed_elements(|store| store.into_iter().map(|(_, value)| value).sum())
367    }
368
369    #[test]
370    fn iter_mut_skips_removed_elements() {
371        iter_skips_removed_elements(|mut store| store.iter_mut().map(|(_, value)| *value).sum())
372    }
373
374    fn iter_skips_removed_elements(sum: impl FnOnce(RwStore<u32>) -> u32) {
375        let store = RwStore::new();
376
377        store.insert(1);
378        let id = store.insert(2);
379        store.insert(4);
380
381        store.remove(id).unwrap();
382
383        assert_eq!(sum(store), 5);
384    }
385
386    #[test]
387    fn into_iter_implements_sync() {
388        let store = RwStore::<*const ()>::new();
389        let iter = store.into_iter();
390        &iter as &dyn Sync;
391    }
392
393    #[test]
394    fn iter_mut_implements_sync() {
395        let mut store = RwStore::<*const ()>::new();
396        let iter = store.iter_mut();
397        &iter as &dyn Sync;
398    }
399
400    #[test]
401    fn into_iter_implements_send() {
402        let store = RwStore::<u32>::new();
403        let iter = store.into_iter();
404        &iter as &dyn Send;
405    }
406
407    #[test]
408    fn iter_mut_implements_send() {
409        let mut store = RwStore::<u32>::new();
410        let iter = store.iter_mut();
411        &iter as &dyn Send;
412    }
413
414    #[test]
415    fn into_iter_implements_unwind_safe() {
416        let store = RwStore::<u32>::new();
417        let iter = store.into_iter();
418        &iter as &dyn UnwindSafe;
419    }
420
421    #[test]
422    fn iter_mut_implements_unwind_safe() {
423        let mut store = RwStore::<u32>::new();
424        let iter = store.iter_mut();
425        &iter as &dyn UnwindSafe;
426    }
427
428    #[test]
429    fn into_iter_implements_ref_unwind_safe() {
430        let store = RwStore::<u32>::new();
431        let iter = store.into_iter();
432        &iter as &dyn RefUnwindSafe;
433    }
434
435    #[test]
436    fn iter_mut_implements_ref_unwind_safe() {
437        let mut store = RwStore::<u32>::new();
438        let iter = store.iter_mut();
439        &iter as &dyn RefUnwindSafe;
440    }
441}