constify from_fn, try_from_fn, try_map, map

Author: bend-n

library/core/src/array/drain.rs

@@ -1,3 +1,4 @@
+#![allow(dead_code)]
 use crate::iter::{TrustedLen, UncheckedIterator};
 use crate::mem::ManuallyDrop;
 use crate::ptr::drop_in_place;

library/core/src/array/mod.rs

@@ -12,7 +12,8 @@ use crate::fmt;
 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,
 };
@@ -24,7 +25,6 @@ mod drain;
 mod equality;
 mod iter;
 
-pub(crate) use drain::drain_array_with;
 #[stable(feature = "array_value_iter", since = "1.51.0")]
 pub use iter::IntoIter;
 
@@ -104,9 +104,10 @@ pub fn repeat<T: Clone, const N: usize>(val: T) -> [T; N] {
 /// ```
 #[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 = "none")]
+pub const fn from_fn<T, 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
 }
@@ -142,11 +143,12 @@ where
 /// ```
 #[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]>,
+    F: [const] FnMut(usize) -> R + [const] Destruct,
+    R: [const] Try<Residual: Residual<[R::Output; N]>>,
+    <R::Residual as Residual<[R::Output; N]>>::TryType: [const] Try,
 {
     let mut array = [const { MaybeUninit::uninit() }; N];
     match try_from_fn_erased(&mut array, cb) {
@@ -542,9 +544,10 @@ impl<T, const N: usize> [T; N] {
     /// ```
     #[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 = "none")]
+    pub const fn map<F, U>(self, f: F) -> [U; N]
     where
-        F: FnMut(T) -> U,
+        F: [const] FnMut(T) -> U + [const] Destruct,
     {
         self.try_map(NeverShortCircuit::wrap_mut_1(f)).0
     }
@@ -580,11 +583,60 @@ impl<T, const N: usize> [T; N] {
     /// 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, F>(self, mut f: F) -> ChangeOutputType<R, [R::Output; N]>
     where
-        R: Try<Residual: Residual<[R::Output; N]>>,
+        F: [const] FnMut(T) -> R + [const] Destruct,
+        R: [const] Try<Residual: Residual<[R::Output; N]>>,
+        <R::Residual as Residual<[R::Output; N]>>::TryType: [const] Try,
     {
-        drain_array_with(self, |iter| try_from_trusted_iterator(iter.map(f)))
+        #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+        struct Guardian<'a, T, U, const N: usize, F: FnMut(T) -> U> {
+            array: ManuallyDrop<[T; N]>,
+            moved: usize,
+            f: &'a mut F,
+        }
+        #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+        impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Guardian<'_, T, U, N, F>
+        where
+            F: [const] FnMut(T) -> U,
+        {
+            type Output = U;
+
+            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")]
+        impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Guardian<'_, T, U, N, F>
+        where
+            F: [const] FnMut(T) -> U,
+        {
+            extern "rust-call" fn call_mut(&mut self, (x,): (usize,)) -> Self::Output {
+                // SAFETY: increment moved before moving. if `f` panics, we drop the rest.
+                self.moved += 1;
+                // SAFETY: caller guarantees never called with number >= N
+                (self.f)(unsafe { self.array.as_ptr().add(x).read() })
+            }
+        }
+        #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+        impl<T: [const] Destruct, U, const N: usize, F: FnMut(T) -> U> const Drop
+            for Guardian<'_, T, U, N, F>
+        {
+            fn drop(&mut self) {
+                let mut n = self.moved;
+                while n != N {
+                    // SAFETY: moved must always be < N
+                    unsafe { self.array.as_mut_ptr().add(n).drop_in_place() };
+                    n += 1;
+                }
+            }
+        }
+        let mut f = Guardian { array: ManuallyDrop::new(self), moved: 0, f: &mut f };
+        // SAFETY: try_from_fn calls `f` with 0..N.
+        let out = try_from_fn(&mut f);
+        mem::forget(f); // it doesnt like being remembered
+        out
     }
 
     /// Returns a slice containing the entire array. Equivalent to `&s[..]`.
@@ -878,12 +930,14 @@ where
 /// 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>(
+#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
+const fn try_from_fn_erased<T, R, F>(
     buffer: &mut [MaybeUninit<T>],
-    mut generator: impl FnMut(usize) -> R,
+    mut generator: F,
 ) -> ControlFlow<R::Residual>
 where
-    R: Try<Output = T>,
+    R: [const] Try<Output = T>,
+    F: [const] FnMut(usize) -> R + [const] Destruct,
 {
     let mut guard = Guard { array_mut: buffer, initialized: 0 };
 
@@ -923,7 +977,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.
@@ -934,15 +989,21 @@ impl<T> Guard<'_, T> {
     }
 }
 
-impl<T> Drop for Guard<'_, T> {
+#[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
+impl<T> 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();
-        }
+        crate::intrinsics::const_eval_select!(
+            @capture [T] { x: &mut Guard<'_, T> = self } -> ():
+            if const {}
+            else {
+                // SAFETY: this slice will contain only initialized objects.
+                unsafe {
+                    x.array_mut.get_unchecked_mut(..x.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.
@@ -380,17 +381,36 @@ pub(crate) type ChangeOutputType<T: Try<Residual: Residual<V>>, V> =
 /// Not currently planned to be exposed publicly, so just `pub(crate)`.
 #[repr(transparent)]
 pub(crate) struct NeverShortCircuit<T>(pub T);
+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`.
     ///
     /// 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]
@@ -401,7 +421,8 @@ impl<T> NeverShortCircuit<T> {
 
 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;
 
@@ -415,15 +436,15 @@ impl<T> Try for NeverShortCircuit<T> {
         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 Residual<T> for NeverShortCircuitResidual {
     type TryType = NeverShortCircuit<T>;
 }
 

library/core/src/ptr/mod.rs

@@ -403,7 +403,7 @@
 
 use crate::cmp::Ordering;
 use crate::intrinsics::const_eval_select;
-use crate::marker::{FnPtr, PointeeSized};
+use crate::marker::{Destruct, FnPtr, PointeeSized};
 use crate::mem::{self, MaybeUninit, SizedTypeProperties};
 use crate::num::NonZero;
 use crate::{fmt, hash, intrinsics, ub_checks};
@@ -801,7 +801,8 @@ pub const unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize) {
 #[lang = "drop_in_place"]
 #[allow(unconditional_recursion)]
 #[rustc_diagnostic_item = "ptr_drop_in_place"]
-pub unsafe fn drop_in_place<T: PointeeSized>(to_drop: *mut T) {
+#[rustc_const_unstable(feature = "const_drop_in_place", issue = "none")]
+pub const unsafe fn drop_in_place<T: PointeeSized + [const] Destruct>(to_drop: *mut T) {
     // Code here does not matter - this is replaced by the
     // real drop glue by the compiler.
 

library/core/src/ptr/mut_ptr.rs

@@ -1390,8 +1390,12 @@ impl<T: PointeeSized> *mut T {
     ///
     /// [`ptr::drop_in_place`]: crate::ptr::drop_in_place()
     #[stable(feature = "pointer_methods", since = "1.26.0")]
+    #[rustc_const_unstable(feature = "const_drop_in_place", issue = "none")]
     #[inline(always)]
-    pub unsafe fn drop_in_place(self) {
+    pub const unsafe fn drop_in_place(self)
+    where
+        T: [const] crate::marker::Destruct,
+    {
         // SAFETY: the caller must uphold the safety contract for `drop_in_place`.
         unsafe { drop_in_place(self) }
     }