|
| 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::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(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 = 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 | +} |
0 commit comments