Rollup merge of #147071 - bend-n:const_array-ops, r=oli-obk

constify from_fn, try_from_fn, try_map, map

adds the `const_array` feature
reimplements `try_map` in more or less the same way

library/core/src/array/drain.rs

@@ -1,76 +1,108 @@
-use crate::iter::{TrustedLen, UncheckedIterator};
-use crate::mem::ManuallyDrop;
-use crate::ptr::drop_in_place;
-use crate::slice;
+use crate::marker::{Destruct, PhantomData};
+use crate::mem::{ManuallyDrop, SizedTypeProperties, conjure_zst};
+use crate::ptr::{NonNull, drop_in_place, from_raw_parts_mut, null_mut};
 
-/// A situationally-optimized version of `array.into_iter().for_each(func)`.
-///
-/// [`crate::array::IntoIter`]s are great when you need an owned iterator, but
-/// storing the entire array *inside* the iterator like that can sometimes
-/// pessimize code.  Notable, it can be more bytes than you really want to move
-/// around, and because the array accesses index into it SRoA has a harder time
-/// optimizing away the type than it does iterators that just hold a couple pointers.
-///
-/// Thus this function exists, which gives a way to get *moved* access to the
-/// elements of an array using a small iterator -- no bigger than a slice iterator.
-///
-/// The function-taking-a-closure structure makes it safe, as it keeps callers
-/// from looking at already-dropped elements.
-pub(crate) fn drain_array_with<T, R, const N: usize>(
-    array: [T; N],
-    func: impl for<'a> FnOnce(Drain<'a, T>) -> R,
-) -> R {
-    let mut array = ManuallyDrop::new(array);
-    // SAFETY: Now that the local won't drop it, it's ok to construct the `Drain` which will.
-    let drain = Drain(array.iter_mut());
-    func(drain)
+impl<'l, 'f, T, U, const N: usize, F: FnMut(T) -> U> Drain<'l, 'f, T, N, F> {
+    /// This function returns a function that lets you index the given array in const.
+    /// As implemented it can optimize better than iterators, and can be constified.
+    /// It acts like a sort of guard (owns the array) and iterator combined, which can be implemented
+    /// as it is a struct that implements const fn;
+    /// in that regard it is somewhat similar to an array::Iter implementing `UncheckedIterator`.
+    /// The only method you're really allowed to call is `next()`,
+    /// anything else is more or less UB, hence this function being unsafe.
+    /// Moved elements will not be dropped.
+    /// This will also not actually store the array.
+    ///
+    /// SAFETY: must only be called `N` times. Thou shalt not drop the array either.
+    // FIXME(const-hack): this is a hack for `let guard = Guard(array); |i| f(guard[i])`.
+    #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+    pub(super) const unsafe fn new(array: &'l mut ManuallyDrop<[T; N]>, f: &'f mut F) -> Self {
+        // dont drop the array, transfers "ownership" to Self
+        let ptr: NonNull<T> = NonNull::from_mut(array).cast();
+        // SAFETY:
+        // Adding `slice.len()` to the starting pointer gives a pointer
+        // at the end of `slice`. `end` will never be dereferenced, only checked
+        // for direct pointer equality with `ptr` to check if the drainer is done.
+        unsafe {
+            let end = if T::IS_ZST { null_mut() } else { ptr.as_ptr().add(N) };
+            Self { ptr, end, f, l: PhantomData }
+        }
+    }
 }
 
-/// See [`drain_array_with`] -- this is `pub(crate)` only so it's allowed to be
-/// mentioned in the signature of that method.  (Otherwise it hits `E0446`.)
-// INVARIANT: It's ok to drop the remainder of the inner iterator.
-pub(crate) struct Drain<'a, T>(slice::IterMut<'a, T>);
+/// See [`Drain::new`]; this is our fake iterator.
+#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+#[unstable(feature = "array_try_map", issue = "79711")]
+pub(super) struct Drain<'l, 'f, T, const N: usize, F> {
+    // FIXME(const-hack): This is essentially a slice::IterMut<'static>, replace when possible.
+    /// The pointer to the next element to return, or the past-the-end location
+    /// if the drainer is empty.
+    ///
+    /// This address will be used for all ZST elements, never changed.
+    /// As we "own" this array, we dont need to store any lifetime.
+    ptr: NonNull<T>,
+    /// For non-ZSTs, the non-null pointer to the past-the-end element.
+    /// For ZSTs, this is null.
+    end: *mut T,
 
-impl<T> Drop for Drain<'_, T> {
+    f: &'f mut F,
+    l: PhantomData<&'l mut [T; N]>,
+}
+
+#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+#[unstable(feature = "array_try_map", issue = "79711")]
+impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, '_, T, N, F>
+where
+    F: [const] FnMut(T) -> U,
+{
+    type Output = U;
+
+    /// This implementation is useless.
+    extern "rust-call" fn call_once(mut self, args: (usize,)) -> Self::Output {
+        self.call_mut(args)
+    }
+}
+#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+#[unstable(feature = "array_try_map", issue = "79711")]
+impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, '_, T, N, F>
+where
+    F: [const] FnMut(T) -> U,
+{
+    // FIXME(const-hack): ideally this would be an unsafe fn `next()`, and to use it you would instead `|_| unsafe { drain.next() }`.
+    extern "rust-call" fn call_mut(
+        &mut self,
+        (_ /* ignore argument */,): (usize,),
+    ) -> Self::Output {
+        if T::IS_ZST {
+            // its UB to call this more than N times, so returning more ZSTs is valid.
+            // SAFETY: its a ZST? we conjur.
+            (self.f)(unsafe { conjure_zst::<T>() })
+        } else {
+            // increment before moving; if `f` panics, we drop the rest.
+            let p = self.ptr;
+            // SAFETY: caller guarantees never called more than N times (see `Drain::new`)
+            self.ptr = unsafe { self.ptr.add(1) };
+            // SAFETY: we are allowed to move this.
+            (self.f)(unsafe { p.read() })
+        }
+    }
+}
+#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+#[unstable(feature = "array_try_map", issue = "79711")]
+impl<T: [const] Destruct, const N: usize, F> const Drop for Drain<'_, '_, T, N, F> {
     fn drop(&mut self) {
-        // SAFETY: By the type invariant, we're allowed to drop all these.
-        unsafe { drop_in_place(self.0.as_mut_slice()) }
-    }
-}
-
-impl<T> Iterator for Drain<'_, T> {
-    type Item = T;
-
-    #[inline]
-    fn next(&mut self) -> Option<T> {
-        let p: *const T = self.0.next()?;
-        // SAFETY: The iterator was already advanced, so we won't drop this later.
-        Some(unsafe { p.read() })
-    }
-
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        let n = self.len();
-        (n, Some(n))
-    }
-}
-
-impl<T> ExactSizeIterator for Drain<'_, T> {
-    #[inline]
-    fn len(&self) -> usize {
-        self.0.len()
-    }
-}
-
-// SAFETY: This is a 1:1 wrapper for a slice iterator, which is also `TrustedLen`.
-unsafe impl<T> TrustedLen for Drain<'_, T> {}
-
-impl<T> UncheckedIterator for Drain<'_, T> {
-    unsafe fn next_unchecked(&mut self) -> T {
-        // SAFETY: `Drain` is 1:1 with the inner iterator, so if the caller promised
-        // that there's an element left, the inner iterator has one too.
-        let p: *const T = unsafe { self.0.next_unchecked() };
-        // SAFETY: The iterator was already advanced, so we won't drop this later.
-        unsafe { p.read() }
+        if !T::IS_ZST {
+            // SAFETY: we cant read more than N elements
+            let slice = unsafe {
+                from_raw_parts_mut::<[T]>(
+                    self.ptr.as_ptr(),
+                    // SAFETY: `start <= end`
+                    self.end.offset_from_unsigned(self.ptr.as_ptr()),
+                )
+            };
+
+            // SAFETY: By the type invariant, we're allowed to drop all these. (we own it, after all)
+            unsafe { drop_in_place(slice) }
+        }
     }
 }

library/core/src/array/mod.rs

@@ -12,7 +12,8 @@
 use crate::hash::{self, Hash};
 use crate::intrinsics::transmute_unchecked;
 use crate::iter::{UncheckedIterator, repeat_n};
-use crate::mem::{self, MaybeUninit};
+use crate::marker::Destruct;
+use crate::mem::{self, ManuallyDrop, MaybeUninit};
 use crate::ops::{
     ChangeOutputType, ControlFlow, FromResidual, Index, IndexMut, NeverShortCircuit, Residual, Try,
 };
@@ -25,7 +26,6 @@
 mod equality;
 mod iter;
 
-pub(crate) use drain::drain_array_with;
 #[stable(feature = "array_value_iter", since = "1.51.0")]
 pub use iter::IntoIter;
 
@@ -105,9 +105,10 @@
 /// ```
 #[inline]
 #[stable(feature = "array_from_fn", since = "1.63.0")]
-pub fn from_fn<T, const N: usize, F>(f: F) -> [T; N]
+#[rustc_const_unstable(feature = "const_array", issue = "147606")]
+pub const fn from_fn<T: [const] Destruct, const N: usize, F>(f: F) -> [T; N]
 where
-    F: FnMut(usize) -> T,
+    F: [const] FnMut(usize) -> T + [const] Destruct,
 {
     try_from_fn(NeverShortCircuit::wrap_mut_1(f)).0
 }
@@ -143,11 +144,11 @@
 /// ```
 #[inline]
 #[unstable(feature = "array_try_from_fn", issue = "89379")]
-pub fn try_from_fn<R, const N: usize, F>(cb: F) -> ChangeOutputType<R, [R::Output; N]>
+#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
+pub const fn try_from_fn<R, const N: usize, F>(cb: F) -> ChangeOutputType<R, [R::Output; N]>
 where
-    F: FnMut(usize) -> R,
-    R: Try,
-    R::Residual: Residual<[R::Output; N]>,
+    R: [const] Try<Residual: [const] Residual<[R::Output; N]>, Output: [const] Destruct>,
+    F: [const] FnMut(usize) -> R + [const] Destruct,
 {
     let mut array = [const { MaybeUninit::uninit() }; N];
     match try_from_fn_erased(&mut array, cb) {
@@ -549,9 +550,12 @@ macro_rules! array_impl_default {
     /// ```
     #[must_use]
     #[stable(feature = "array_map", since = "1.55.0")]
-    pub fn map<F, U>(self, f: F) -> [U; N]
+    #[rustc_const_unstable(feature = "const_array", issue = "147606")]
+    pub const fn map<F, U>(self, f: F) -> [U; N]
     where
-        F: FnMut(T) -> U,
+        F: [const] FnMut(T) -> U + [const] Destruct,
+        U: [const] Destruct,
+        T: [const] Destruct,
     {
         self.try_map(NeverShortCircuit::wrap_mut_1(f)).0
     }
@@ -587,11 +591,19 @@ macro_rules! array_impl_default {
     /// assert_eq!(c, Some(a));
     /// ```
     #[unstable(feature = "array_try_map", issue = "79711")]
-    pub fn try_map<R>(self, f: impl FnMut(T) -> R) -> ChangeOutputType<R, [R::Output; N]>
+    #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+    pub const fn try_map<R>(
+        self,
+        mut f: impl [const] FnMut(T) -> R + [const] Destruct,
+    ) -> ChangeOutputType<R, [R::Output; N]>
     where
-        R: Try<Residual: Residual<[R::Output; N]>>,
+        R: [const] Try<Residual: [const] Residual<[R::Output; N]>, Output: [const] Destruct>,
+        T: [const] Destruct,
     {
-        drain_array_with(self, |iter| try_from_trusted_iterator(iter.map(f)))
+        let mut me = ManuallyDrop::new(self);
+        // SAFETY: try_from_fn calls `f` N times.
+        let mut f = unsafe { drain::Drain::new(&mut me, &mut f) };
+        try_from_fn(&mut f)
     }
 
     /// Returns a slice containing the entire array. Equivalent to `&s[..]`.
@@ -885,13 +897,11 @@ fn next<T>(mut iter: impl UncheckedIterator<Item = T>) -> impl FnMut(usize) -> T
 /// not optimizing away.  So if you give it a shot, make sure to watch what
 /// happens in the codegen tests.
 #[inline]
-fn try_from_fn_erased<T, R>(
-    buffer: &mut [MaybeUninit<T>],
-    mut generator: impl FnMut(usize) -> R,
-) -> ControlFlow<R::Residual>
-where
-    R: Try<Output = T>,
-{
+#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
+const fn try_from_fn_erased<R: [const] Try<Output: [const] Destruct>>(
+    buffer: &mut [MaybeUninit<R::Output>],
+    mut generator: impl [const] FnMut(usize) -> R + [const] Destruct,
+) -> ControlFlow<R::Residual> {
     let mut guard = Guard { array_mut: buffer, initialized: 0 };
 
     while guard.initialized < guard.array_mut.len() {
@@ -930,7 +940,8 @@ impl<T> Guard<'_, T> {
     ///
     /// No more than N elements must be initialized.
     #[inline]
-    pub(crate) unsafe fn push_unchecked(&mut self, item: T) {
+    #[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
+    pub(crate) const unsafe fn push_unchecked(&mut self, item: T) {
         // SAFETY: If `initialized` was correct before and the caller does not
         // invoke this method more than N times then writes will be in-bounds
         // and slots will not be initialized more than once.
@@ -941,11 +952,11 @@ pub(crate) unsafe fn push_unchecked(&mut self, item: T) {
     }
 }
 
-impl<T> Drop for Guard<'_, T> {
+#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
+impl<T: [const] Destruct> const Drop for Guard<'_, T> {
     #[inline]
     fn drop(&mut self) {
         debug_assert!(self.initialized <= self.array_mut.len());
-
         // SAFETY: this slice will contain only initialized objects.
         unsafe {
             self.array_mut.get_unchecked_mut(..self.initialized).assume_init_drop();

library/core/src/ops/try_trait.rs

@@ -1,3 +1,4 @@
+use crate::marker::{Destruct, PhantomData};
 use crate::ops::ControlFlow;
 
 /// The `?` operator and `try {}` blocks.
@@ -363,6 +364,7 @@ pub fn from_yeet<T, Y>(yeeted: Y) -> T
 pub const trait Residual<O>: Sized {
     /// The "return" type of this meta-function.
     #[unstable(feature = "try_trait_v2_residual", issue = "91285")]
+    // FIXME: ought to be implied
     type TryType: [const] Try<Output = O, Residual = Self>;
 }
 
@@ -396,6 +398,25 @@ pub const fn residual_into_try_type<R: [const] Residual<O>, O>(
 /// Not currently planned to be exposed publicly, so just `pub(crate)`.
 #[repr(transparent)]
 pub(crate) struct NeverShortCircuit<T>(pub T);
+// FIXME(const-hack): replace with `|a| NeverShortCircuit(f(a))` when const closures added.
+pub(crate) struct Wrapped<T, A, F: FnMut(A) -> T> {
+    f: F,
+    p: PhantomData<(T, A)>,
+}
+#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
+impl<T, A, F: [const] FnMut(A) -> T + [const] Destruct> const FnOnce<(A,)> for Wrapped<T, A, F> {
+    type Output = NeverShortCircuit<T>;
+
+    extern "rust-call" fn call_once(mut self, args: (A,)) -> Self::Output {
+        self.call_mut(args)
+    }
+}
+#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
+impl<T, A, F: [const] FnMut(A) -> T> const FnMut<(A,)> for Wrapped<T, A, F> {
+    extern "rust-call" fn call_mut(&mut self, (args,): (A,)) -> Self::Output {
+        NeverShortCircuit((self.f)(args))
+    }
+}
 
 impl<T> NeverShortCircuit<T> {
     /// Wraps a unary function to produce one that wraps the output into a `NeverShortCircuit`.
@@ -403,10 +424,11 @@ impl<T> NeverShortCircuit<T> {
     /// This is useful for implementing infallible functions in terms of the `try_` ones,
     /// without accidentally capturing extra generic parameters in a closure.
     #[inline]
-    pub(crate) fn wrap_mut_1<A>(
-        mut f: impl FnMut(A) -> T,
-    ) -> impl FnMut(A) -> NeverShortCircuit<T> {
-        move |a| NeverShortCircuit(f(a))
+    pub(crate) const fn wrap_mut_1<A, F>(f: F) -> Wrapped<T, A, F>
+    where
+        F: [const] FnMut(A) -> T,
+    {
+        Wrapped { f, p: PhantomData }
     }
 
     #[inline]
@@ -417,7 +439,8 @@ pub(crate) fn wrap_mut_2<A, B>(mut f: impl FnMut(A, B) -> T) -> impl FnMut(A, B)
 
 pub(crate) enum NeverShortCircuitResidual {}
 
-impl<T> Try for NeverShortCircuit<T> {
+#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
+impl<T> const Try for NeverShortCircuit<T> {
     type Output = T;
     type Residual = NeverShortCircuitResidual;
 
@@ -431,15 +454,15 @@ fn from_output(x: T) -> Self {
         NeverShortCircuit(x)
     }
 }
-
-impl<T> FromResidual for NeverShortCircuit<T> {
+#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
+impl<T> const FromResidual for NeverShortCircuit<T> {
     #[inline]
     fn from_residual(never: NeverShortCircuitResidual) -> Self {
         match never {}
     }
 }
-
-impl<T> Residual<T> for NeverShortCircuitResidual {
+#[rustc_const_unstable(feature = "const_never_short_circuit", issue = "none")]
+impl<T: [const] Destruct> const Residual<T> for NeverShortCircuitResidual {
     type TryType = NeverShortCircuit<T>;
 }
 

library/coretests/tests/array.rs

@@ -724,3 +724,20 @@ fn array_eq() {
     let not_true = [0u8] == [].as_slice();
     assert!(!not_true);
 }
+
+#[test]
+fn const_array_ops() {
+    const fn doubler(x: usize) -> usize {
+        x * 2
+    }
+    const fn maybe_doubler(x: usize) -> Option<usize> {
+        x.checked_mul(2)
+    }
+    assert_eq!(const { std::array::from_fn::<_, 5, _>(doubler) }, [0, 2, 4, 6, 8]);
+    assert_eq!(const { [5, 6, 1, 2].map(doubler) }, [10, 12, 2, 4]);
+    assert_eq!(const { [1, usize::MAX, 2, 8].try_map(maybe_doubler) }, None);
+    assert_eq!(const { std::array::try_from_fn::<_, 5, _>(maybe_doubler) }, Some([0, 2, 4, 6, 8]));
+    #[derive(Debug, PartialEq)]
+    struct Zst;
+    assert_eq!([(); 10].try_map(|()| Some(Zst)), Some([const { Zst }; 10]));
+}

library/coretests/tests/lib.rs

@@ -4,6 +4,7 @@
 #![feature(alloc_layout_extra)]
 #![feature(array_ptr_get)]
 #![feature(array_try_from_fn)]
+#![feature(array_try_map)]
 #![feature(array_windows)]
 #![feature(ascii_char)]
 #![feature(ascii_char_variants)]
@@ -16,6 +17,7 @@
 #![feature(char_internals)]
 #![feature(char_max_len)]
 #![feature(clone_to_uninit)]
+#![feature(const_array)]
 #![feature(const_cell_traits)]
 #![feature(const_cmp)]
 #![feature(const_convert)]