constify from_fn, try_from_fn, try_map, map

Author: bend-n

library/core/src/array/drain.rs

@@ -1,76 +1,61 @@
-use crate::iter::{TrustedLen, UncheckedIterator};
+use crate::assert_unsafe_precondition;
+use crate::marker::Destruct;
 use crate::mem::ManuallyDrop;
-use crate::ptr::drop_in_place;
-use crate::slice;
 
-/// 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)
+#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+#[unstable(feature = "array_try_map", issue = "79711")]
+pub(super) struct Drain<'a, T, U, const N: usize, F: FnMut(T) -> U> {
+ array: ManuallyDrop<[T; N]>,
+ moved: usize,
+ f: &'a mut F,
 }
-
-/// 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>);
-
-impl<T> Drop for Drain<'_, T> {
- 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()) }
+#[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, 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)
  }
 }
-
-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))
+#[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, U, N, F>
+where
+ F: [const] FnMut(T) -> U,
+{
+ extern "rust-call" fn call_mut(&mut self, (i,): (usize,)) -> Self::Output {
+ // SAFETY: increment moved before moving. if `f` panics, we drop the rest.
+ self.moved += 1;
+  assert_unsafe_precondition!(
+ check_library_ub,
+  "musnt index array out of bounds", (i: usize = i, size: usize = N) => i < size
+  );
+ // SAFETY: caller guarantees never called with number >= N (see `Drain::new`)
+ (self.f)(unsafe { self.array.as_ptr().add(i).read() })
  }
 }
-
-impl<T> ExactSizeIterator for Drain<'_, T> {
- #[inline]
- fn len(&self) -> usize {
- self.0.len()
+#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
+#[unstable(feature = "array_try_map", issue = "79711")]
+impl<T: [const] Destruct, U, const N: usize, F: FnMut(T) -> U> const Drop
+ for Drain<'_, 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;
+ }
  }
 }
-
-// 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() }
+impl<'a, T, U, const N: usize, F: FnMut(T) -> U> Drain<'a, T, U, N, F> {
+ /// SAFETY: must be called without indexing out of bounds. (see `Drain::call_mut`)
+ // FIXME(const-hack): this is a hack for `let guard = Guard(array); |i| f(guard[i])`.
+ pub(super) const unsafe fn new(array: [T; N], f: &'a mut F) -> Self {
+ Self { array: ManuallyDrop::new(array), moved: 0, f }
  }
 }

library/core/src/array/mod.rs

@@ -12,6 +12,7 @@ use crate::fmt;
 use crate::hash::{self, Hash};
 use crate::intrinsics::transmute_unchecked;
 use crate::iter::{UncheckedIterator, repeat_n};
+use crate::marker::Destruct;
 use crate::mem::{self, 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 = "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
 }
@@ -142,11 +143,13 @@ 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::Output: [const] Destruct,
+ <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 +545,11 @@ 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 = "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,
  {
  self.try_map(NeverShortCircuit::wrap_mut_1(f)).0
  }
@@ -580,11 +585,19 @@ 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::Output: [const] Destruct,
+ <R::Residual as Residual<[R::Output; N]>>::TryType: [const] Try,
  {
- drain_array_with(self, |iter| try_from_trusted_iterator(iter.map(f)))
+ // SAFETY: try_from_fn calls `f` with 0..N.
+ let mut f = unsafe { drain::Drain::new(self, &mut f) };
+ 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 +891,15 @@ 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>,
+ T: [const] Destruct,
+ R: [const] Try<Output = T>,
+ F: [const] FnMut(usize) -> R + [const] Destruct,
 {
  let mut guard = Guard { array_mut: buffer, initialized: 0 };
 
@@ -923,7 +939,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,11 +951,11 @@ impl<T> Guard<'_, 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.
@@ -380,17 +381,37 @@ 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);
+// 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`.
  ///
  /// 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 +422,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 +437,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>;
 }