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.
drain_array_with<T, R, const N: usize>( array: [T; N], func: impl for<'a> FnOnce(Drain<'a, T>) -> R, ) -> R19 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> {
drop(&mut self)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]
next(&mut self) -> Option<T>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]
size_hint(&self) -> (usize, Option<usize>)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]
len(&self) -> usize60 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> {
next_unchecked(&mut self) -> T69 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