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#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
35#[unstable(feature = "array_try_map", issue = "79711")]
36pub(super) struct Drain<'l, 'f, T, const N: usize, F> {
37    // FIXME(const-hack): This is essentially a slice::IterMut<'static>, replace when possible.
38    /// The pointer to the next element to return, or the past-the-end location
39    /// if the drainer is empty.
40    ///
41    /// This address will be used for all ZST elements, never changed.
42    /// As we "own" this array, we dont need to store any lifetime.
43    ptr: NonNull<T>,
44    /// For non-ZSTs, the non-null pointer to the past-the-end element.
45    /// For ZSTs, this is null.
46    end: *mut T,
47
48    f: &'f mut F,
49    l: PhantomData<&'l mut [T; N]>,
50}
51
52#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
53#[unstable(feature = "array_try_map", issue = "79711")]
54impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, '_, T, N, F>
55where
56    F: [const] FnMut(T) -> U,
57{
58    type Output = U;
59
60    /// This implementation is useless.
61    extern "rust-call" fn call_once(mut self, args: (usize,)) -> Self::Output {
62        self.call_mut(args)
63    }
64}
65#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
66#[unstable(feature = "array_try_map", issue = "79711")]
67impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, '_, T, N, F>
68where
69    F: [const] FnMut(T) -> U,
70{
71    // FIXME(const-hack): ideally this would be an unsafe fn `next()`, and to use it you would instead `|_| unsafe { drain.next() }`.
72    extern "rust-call" fn call_mut(
73        &mut self,
74        (_ /* ignore argument */,): (usize,),
75    ) -> Self::Output {
76        if T::IS_ZST {
77            // its UB to call this more than N times, so returning more ZSTs is valid.
78            // SAFETY: its a ZST? we conjur.
79            (self.f)(unsafe { conjure_zst::<T>() })
80        } else {
81            // increment before moving; if `f` panics, we drop the rest.
82            let p = self.ptr;
83            // SAFETY: caller guarantees never called more than N times (see `Drain::new`)
84            self.ptr = unsafe { self.ptr.add(1) };
85            // SAFETY: we are allowed to move this.
86            (self.f)(unsafe { p.read() })
87        }
88    }
89}
90#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
91#[unstable(feature = "array_try_map", issue = "79711")]
92impl<T: [const] Destruct, const N: usize, F> const Drop for Drain<'_, '_, T, N, F> {
93    fn drop(&mut self) {
94        if !T::IS_ZST {
95            // SAFETY: we cant read more than N elements
96            let slice = unsafe {
97                from_raw_parts_mut::<[T]>(
98                    self.ptr.as_ptr(),
99                    // SAFETY: `start <= end`
100                    self.end.offset_from_unsigned(self.ptr.as_ptr()),
101                )
102            };
103
104            // SAFETY: By the type invariant, we're allowed to drop all these. (we own it, after all)
105            unsafe { drop_in_place(slice) }
106        }
107    }
108}
core/array/drain.rs

core/array/
drain.rs
