1use 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
13pub struct IntoIter<Element> {
26 iter: GenIter<OwnedBlockTracker<Element>>,
27}
28
29impl<Element> IntoIter<Element> {
30 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
71pub struct IterMut<'a, Element> {
84 iter: GenIter<MutBlockTracker<'a, Element>>,
85}
86
87impl<'a, Element> IterMut<'a, Element> {
88 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}