1 | use crate::iter::{TrustedLen, UncheckedIterator}; |
2 | use crate::mem::ManuallyDrop; |
3 | use crate::ptr::drop_in_place; |
4 | use crate::slice; |
5 | |
6 | /// A situationally-optimized version of `array.into_iter().for_each(func)`. |
7 | /// |
8 | /// [`crate::array::IntoIter`]s are great when you need an owned iterator, but |
9 | /// storing the entire array *inside* the iterator like that can sometimes |
10 | /// pessimize code. Notable, it can be more bytes than you really want to move |
11 | /// around, and because the array accesses index into it SRoA has a harder time |
12 | /// optimizing away the type than it does iterators that just hold a couple pointers. |
13 | /// |
14 | /// Thus this function exists, which gives a way to get *moved* access to the |
15 | /// elements of an array using a small iterator -- no bigger than a slice iterator. |
16 | /// |
17 | /// The function-taking-a-closure structure makes it safe, as it keeps callers |
18 | /// from looking at already-dropped elements. |
19 | pub(crate) fn drain_array_with<T, R, const N: usize>( |
20 | array: [T; N], |
21 | func: impl for<'a> FnOnce(Drain<'a, T>) -> R, |
22 | ) -> R { |
23 | let mut array: ManuallyDrop<[T; N]> = ManuallyDrop::new(array); |
24 | // SAFETY: Now that the local won't drop it, it's ok to construct the `Drain` which will. |
25 | let drain: Drain<'_, T> = Drain(array.iter_mut()); |
26 | func(drain) |
27 | } |
28 | |
29 | /// See [`drain_array_with`] -- this is `pub(crate)` only so it's allowed to be |
30 | /// mentioned in the signature of that method. (Otherwise it hits `E0446`.) |
31 | // INVARIANT: It's ok to drop the remainder of the inner iterator. |
32 | pub(crate) struct Drain<'a, T>(slice::IterMut<'a, T>); |
33 | |
34 | impl<T> Drop for Drain<'_, T> { |
35 | fn drop(&mut self) { |
36 | // SAFETY: By the type invariant, we're allowed to drop all these. |
37 | unsafe { drop_in_place(self.0.as_mut_slice()) } |
38 | } |
39 | } |
40 | |
41 | impl<T> Iterator for Drain<'_, T> { |
42 | type Item = T; |
43 | |
44 | #[inline ] |
45 | fn next(&mut self) -> Option<T> { |
46 | let p: *const T = self.0.next()?; |
47 | // SAFETY: The iterator was already advanced, so we won't drop this later. |
48 | Some(unsafe { p.read() }) |
49 | } |
50 | |
51 | #[inline ] |
52 | fn size_hint(&self) -> (usize, Option<usize>) { |
53 | let n: usize = self.len(); |
54 | (n, Some(n)) |
55 | } |
56 | } |
57 | |
58 | impl<T> ExactSizeIterator for Drain<'_, T> { |
59 | #[inline ] |
60 | fn len(&self) -> usize { |
61 | self.0.len() |
62 | } |
63 | } |
64 | |
65 | // SAFETY: This is a 1:1 wrapper for a slice iterator, which is also `TrustedLen`. |
66 | unsafe impl<T> TrustedLen for Drain<'_, T> {} |
67 | |
68 | impl<T> UncheckedIterator for Drain<'_, T> { |
69 | unsafe fn next_unchecked(&mut self) -> T { |
70 | // SAFETY: `Drain` is 1:1 with the inner iterator, so if the caller promised |
71 | // that there's an element left, the inner iterator has one too. |
72 | let p: *const T = unsafe { self.0.next_unchecked() }; |
73 | // SAFETY: The iterator was already advanced, so we won't drop this later. |
74 | unsafe { p.read() } |
75 | } |
76 | } |
77 | |