1use crate::iter::{TrustedLen, UncheckedIterator};
2use crate::mem::ManuallyDrop;
3use crate::ptr::drop_in_place;
4use 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.
19pub(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.
32pub(crate) struct Drain<'a, T>(slice::IterMut<'a, T>);
33
34impl<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
41impl<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
58impl<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`.
66unsafe impl<T> TrustedLen for Drain<'_, T> {}
67
68impl<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