redo the drain

library/core/src/array/drain.rs

@@ -1,76 +1,108 @@
-use crate::assert_unsafe_precondition;
-use crate::marker::Destruct;
-use crate::mem::ManuallyDrop;
+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};
+
+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::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,
+
+    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")]
-pub(super) struct Drain<'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")]
-#[unstable(feature = "array_try_map", issue = "79711")]
-impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, T, U, N, F>
+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, U, N, F>
+impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, '_, T, 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: the `i` should also always go up, and musnt skip any, else some things will be leaked.
-        // SAFETY: if it goes down, we will drop freed elements. not good.
-        // SAFETY: caller guarantees never called with number >= N (see `Drain::new`)
-        (self.f)(unsafe { self.array.as_ptr().add(i).read() })
+    // 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, U, const N: usize, F: FnMut(T) -> U> const Drop
-    for Drain<'_, T, U, N, F>
-{
+impl<T: [const] Destruct, const N: usize, F> const Drop for Drain<'_, '_, T, 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;
+        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) }
         }
     }
 }
-impl<'a, T, U, const N: usize, F: FnMut(T) -> U> Drain<'a, T, U, 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 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.
-    ///
-    /// Previously this was implemented as a wrapper around a `slice::Iter`, which
-    /// called `read()` on the returned `&T`; gnarly stuff.
-    ///
-    /// SAFETY: must be called in order of 0..N, without indexing out of bounds. (see `Drain::call_mut`)
-    /// Potentially the function could completely disregard the supplied argument, however i think that behaviour would be unintuitive.
-    // 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

@@ -13,7 +13,7 @@ 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::mem::{self, ManuallyDrop, MaybeUninit};
 use crate::ops::{
     ChangeOutputType, ControlFlow, FromResidual, Index, IndexMut, NeverShortCircuit, Residual, Try,
 };
@@ -145,14 +145,10 @@ where
 #[inline]
 #[unstable(feature = "array_try_from_fn", issue = "89379")]
 #[rustc_const_unstable(feature = "array_try_from_fn", issue = "89379")]
-pub const fn try_from_fn<R, const N: usize>(
-    cb: impl [const] FnMut(usize) -> R + [const] Destruct,
-) -> ChangeOutputType<R, [R::Output; N]>
+pub const fn try_from_fn<R, const N: usize, F>(cb: F) -> ChangeOutputType<R, [R::Output; N]>
 where
-    R: [const] Try<
-            Residual: [const] Residual<[R::Output; N], TryType: [const] Try>,
-            Output: [const] Destruct,
-        >,
+    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) {
@@ -559,6 +555,7 @@ impl<T, const N: usize> [T; N] {
     where
         F: [const] FnMut(T) -> U + [const] Destruct,
         U: [const] Destruct,
+        T: [const] Destruct,
     {
         self.try_map(NeverShortCircuit::wrap_mut_1(f)).0
     }
@@ -600,16 +597,13 @@ impl<T, const N: usize> [T; N] {
         mut f: impl [const] FnMut(T) -> R + [const] Destruct,
     ) -> ChangeOutputType<R, [R::Output; N]>
     where
-        R: [const] Try<
-                Residual: [const] Residual<[R::Output; N], TryType: [const] Try>,
-                Output: [const] Destruct,
-            >,
+        R: [const] Try<Residual: [const] Residual<[R::Output; N]>, Output: [const] Destruct>,
+        T: [const] Destruct,
     {
-        // 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
+        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[..]`.

library/core/src/ops/try_trait.rs

@@ -364,6 +364,7 @@ where
 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>;
 }
 

library/coretests/tests/array.rs

@@ -737,4 +737,7 @@ fn const_array_ops() {
     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]));
 }