core/array/
drain.rs

1use crate::marker::{Destruct, PhantomData};
2use crate::mem::{ManuallyDrop, SizedTypeProperties, conjure_zst};
3use crate::ptr::{NonNull, drop_in_place, from_raw_parts_mut, null_mut};
4
5impl<'l, 'f, T, U, const N: usize, F: FnMut(T) -> U> Drain<'l, 'f, T, N, F> {
6    /// This function returns a function that lets you index the given array in const.
7    /// As implemented it can optimize better than iterators, and can be constified.
8    /// It acts like a sort of guard (owns the array) and iterator combined, which can be implemented
9    /// as it is a struct that implements const fn;
10    /// in that regard it is somewhat similar to an array::Iter implementing `UncheckedIterator`.
11    /// The only method you're really allowed to call is `next()`,
12    /// anything else is more or less UB, hence this function being unsafe.
13    /// Moved elements will not be dropped.
14    /// This will also not actually store the array.
15    ///
16    /// SAFETY: must only be called `N` times. Thou shalt not drop the array either.
17    // FIXME(const-hack): this is a hack for `let guard = Guard(array); |i| f(guard[i])`.
18    #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
19    pub(super) const unsafe fn new(array: &'l mut ManuallyDrop<[T; N]>, f: &'f mut F) -> Self {
20        // dont drop the array, transfers "ownership" to Self
21        let ptr: NonNull<T> = NonNull::from_mut(array).cast();
22        // SAFETY:
23        // Adding `slice.len()` to the starting pointer gives a pointer
24        // at the end of `slice`. `end` will never be dereferenced, only checked
25        // for direct pointer equality with `ptr` to check if the drainer is done.
26        unsafe {
27            let end = if T::IS_ZST { null_mut() } else { ptr.as_ptr().add(N) };
28            Self { ptr, end, f, l: PhantomData }
29        }
30    }
31}
32
33/// See [`Drain::new`]; this is our fake iterator.
34#[unstable(feature = "array_try_map", issue = "79711")]
35pub(super) struct Drain<'l, 'f, T, const N: usize, F> {
36    // FIXME(const-hack): This is essentially a slice::IterMut<'static>, replace when possible.
37    /// The pointer to the next element to return, or the past-the-end location
38    /// if the drainer is empty.
39    ///
40    /// This address will be used for all ZST elements, never changed.
41    /// As we "own" this array, we dont need to store any lifetime.
42    ptr: NonNull<T>,
43    /// For non-ZSTs, the non-null pointer to the past-the-end element.
44    /// For ZSTs, this is null.
45    end: *mut T,
46
47    f: &'f mut F,
48    l: PhantomData<&'l mut [T; N]>,
49}
50
51#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
52#[unstable(feature = "array_try_map", issue = "79711")]
53impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, '_, T, N, F>
54where
55    F: [const] FnMut(T) -> U,
56{
57    type Output = U;
58
59    /// This implementation is useless.
60    extern "rust-call" fn call_once(mut self, args: (usize,)) -> Self::Output {
61        self.call_mut(args)
62    }
63}
64#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
65#[unstable(feature = "array_try_map", issue = "79711")]
66impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, '_, T, N, F>
67where
68    F: [const] FnMut(T) -> U,
69{
70    // FIXME(const-hack): ideally this would be an unsafe fn `next()`, and to use it you would instead `|_| unsafe { drain.next() }`.
71    extern "rust-call" fn call_mut(
72        &mut self,
73        (_ /* ignore argument */,): (usize,),
74    ) -> Self::Output {
75        if T::IS_ZST {
76            // its UB to call this more than N times, so returning more ZSTs is valid.
77            // SAFETY: its a ZST? we conjur.
78            (self.f)(unsafe { conjure_zst::<T>() })
79        } else {
80            // increment before moving; if `f` panics, we drop the rest.
81            let p = self.ptr;
82            // SAFETY: caller guarantees never called more than N times (see `Drain::new`)
83            self.ptr = unsafe { self.ptr.add(1) };
84            // SAFETY: we are allowed to move this.
85            (self.f)(unsafe { p.read() })
86        }
87    }
88}
89#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
90#[unstable(feature = "array_try_map", issue = "79711")]
91impl<T: [const] Destruct, const N: usize, F> const Drop for Drain<'_, '_, T, N, F> {
92    fn drop(&mut self) {
93        if !T::IS_ZST {
94            // SAFETY: we cant read more than N elements
95            let slice = unsafe {
96                from_raw_parts_mut::<[T]>(
97                    self.ptr.as_ptr(),
98                    // SAFETY: `start <= end`
99                    self.end.offset_from_unsigned(self.ptr.as_ptr()),
100                )
101            };
102
103            // SAFETY: By the type invariant, we're allowed to drop all these. (we own it, after all)
104            unsafe { drop_in_place(slice) }
105        }
106    }
107}
core/array/drain.rs

core/array/
drain.rs
