drain.rs - source
Description: Source of the Rust file `library/core/src/array/drain.rs`.

#
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}