core/array/
drain.rs

1use crate::marker::{Destruct, PhantomData};
2use crate::mem::{ManuallyDrop, SizedTypeProperties, conjure_zst, transmute};
3use crate::ptr::{NonNull, drop_in_place, from_raw_parts_mut, without_provenance_mut};
4
5impl<'l, 'f, T, U, F: FnMut(T) -> U> Drain<'l, 'f, T, 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    /// The only method you're really allowed to call is `next()`,
11    /// anything else is more or less UB, hence this function being unsafe.
12    /// Moved elements will not be dropped.
13    /// This will also not actually store the array.
14    ///
15    /// SAFETY: must only be called `N` times. Thou shalt not drop the array either.
16    #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
17    pub(super) const unsafe fn new<const N: usize>(
18        array: &'l mut ManuallyDrop<[T; N]>,
19        f: &'f mut F,
20    ) -> Self {
21        // dont drop the array, transfers "ownership" to Self
22        let ptr: NonNull<T> = NonNull::from_mut(array).cast();
23        // SAFETY:
24        // Adding `slice.len()` to the starting pointer gives a pointer
25        // at the end of `slice`. `end` will never be dereferenced, only checked
26        // for direct pointer equality with `ptr` to check if the drainer is done.
27        unsafe {
28            let end_or_len =
29                if T::IS_ZST { without_provenance_mut(N) } else { ptr.as_ptr().add(N) };
30            Self { ptr, end_or_len, f, l: PhantomData }
31        }
32    }
33}
34
35/// See [`Drain::new`]; this is our fake iterator.
36#[unstable(feature = "array_try_map", issue = "79711")]
37pub(super) struct Drain<'l, 'f, T, F> {
38    // FIXME(const-hack): This is a slice::IterMut<'l>, replace when possible.
39    /// The pointer to the next element to return, or the past-the-end location
40    /// if the drainer is empty.
41    ///
42    /// This address will be used for all ZST elements, never changed.
43    /// As we "own" this array, we dont need to store any lifetime.
44    ptr: NonNull<T>,
45    /// For non-ZSTs, the non-null pointer to the past-the-end element.
46    /// For ZSTs, this is the number of unprocessed items.
47    end_or_len: *mut T,
48
49    f: &'f mut F,
50    l: PhantomData<&'l mut [T]>,
51}
52
53#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
54#[unstable(feature = "array_try_map", issue = "79711")]
55impl<T, U, F> const FnOnce<(usize,)> for &mut Drain<'_, '_, T, F>
56where
57    F: [const] FnMut(T) -> U,
58{
59    type Output = U;
60
61    /// This implementation is useless.
62    extern "rust-call" fn call_once(mut self, args: (usize,)) -> Self::Output {
63        self.call_mut(args)
64    }
65}
66#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
67#[unstable(feature = "array_try_map", issue = "79711")]
68impl<T, U, F> const FnMut<(usize,)> for &mut Drain<'_, '_, T, F>
69where
70    F: [const] FnMut(T) -> U,
71{
72    // FIXME(const-hack): ideally this would be an unsafe fn `next()`, and to use it you would instead `|_| unsafe { drain.next() }`.
73    extern "rust-call" fn call_mut(
74        &mut self,
75        (_ /* ignore argument */,): (usize,),
76    ) -> Self::Output {
77        if T::IS_ZST {
78            #[expect(ptr_to_integer_transmute_in_consts)]
79            // SAFETY:
80            // This is equivalent to `self.end_or_len.addr`, but that's not
81            // available in `const`. `self.end_or_len` doesn't have provenance,
82            // so transmuting is fine.
83            let len = unsafe { transmute::<*mut T, usize>(self.end_or_len) };
84            // SAFETY:
85            // The caller guarantees that this is never called more than N times
86            // (see `Drain::new`), hence this cannot underflow.
87            self.end_or_len = without_provenance_mut(unsafe { len.unchecked_sub(1) });
88            // its UB to call this more than N times, so returning more ZSTs is valid.
89            // SAFETY: its a ZST? we conjur.
90            (self.f)(unsafe { conjure_zst::<T>() })
91        } else {
92            // increment before moving; if `f` panics, we drop the rest.
93            let p = self.ptr;
94            // SAFETY: caller guarantees never called more than N times (see `Drain::new`)
95            self.ptr = unsafe { self.ptr.add(1) };
96            // SAFETY: we are allowed to move this.
97            (self.f)(unsafe { p.read() })
98        }
99    }
100}
101#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
102#[unstable(feature = "array_try_map", issue = "79711")]
103impl<T: [const] Destruct, F> const Drop for Drain<'_, '_, T, F> {
104    fn drop(&mut self) {
105        let slice = if T::IS_ZST {
106            from_raw_parts_mut::<[T]>(
107                self.ptr.as_ptr(),
108                #[expect(ptr_to_integer_transmute_in_consts)]
109                // SAFETY:
110                // This is equivalent to `self.end_or_len.addr`, but that's not
111                // available in `const`. `self.end_or_len` doesn't have provenance,
112                // so transmuting is fine.
113                unsafe {
114                    transmute::<*mut T, usize>(self.end_or_len)
115                },
116            )
117        } else {
118            // SAFETY: we cant read more than N elements
119            unsafe {
120                from_raw_parts_mut::<[T]>(
121                    self.ptr.as_ptr(),
122                    // SAFETY: `start <= end`
123                    self.end_or_len.offset_from_unsigned(self.ptr.as_ptr()),
124                )
125            }
126        };
127
128        // SAFETY: By the type invariant, we're allowed to drop all these. (we own it, after all)
129        unsafe { drop_in_place(slice) }
130    }
131}
core/array/drain.rs

core/array/
drain.rs
