From e3a2c23e37f16e08b316cd83c9ee8875d826e9ff Mon Sep 17 00:00:00 2001 From: bendn Date: Mon, 24 Nov 2025 17:09:41 +0700 Subject: [PATCH] redo the drain --- library/core/src/array/drain.rs | 134 ++++++++++++++++++------------ library/core/src/array/mod.rs | 28 +++---- library/core/src/ops/try_trait.rs | 1 + library/coretests/tests/array.rs | 3 + 4 files changed, 98 insertions(+), 68 deletions(-) diff --git a/library/core/src/array/drain.rs b/library/core/src/array/drain.rs index 6ab649c036b4..1c6137191324 100644 --- a/library/core/src/array/drain.rs +++ b/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 = 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, + /// 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 const FnOnce<(usize,)> for &mut Drain<'_, T, U, N, F> +impl 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 const FnMut<(usize,)> for &mut Drain<'_, T, U, N, F> +impl 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::() }) + } 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 U> const Drop - for Drain<'_, T, U, N, F> -{ +impl 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 } - } -} diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index bbf75c70d696..c8d0622059cd 100644 --- a/library/core/src/array/mod.rs +++ b/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( - cb: impl [const] FnMut(usize) -> R + [const] Destruct, -) -> ChangeOutputType +pub const fn try_from_fn(cb: F) -> ChangeOutputType where - R: [const] Try< - Residual: [const] Residual<[R::Output; N], TryType: [const] Try>, - Output: [const] Destruct, - >, + R: [const] Try, 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; 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; N] { mut f: impl [const] FnMut(T) -> R + [const] Destruct, ) -> ChangeOutputType where - R: [const] Try< - Residual: [const] Residual<[R::Output; N], TryType: [const] Try>, - Output: [const] Destruct, - >, + R: [const] Try, 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[..]`. diff --git a/library/core/src/ops/try_trait.rs b/library/core/src/ops/try_trait.rs index bcff1d7f456c..34000f6d6b21 100644 --- a/library/core/src/ops/try_trait.rs +++ b/library/core/src/ops/try_trait.rs @@ -364,6 +364,7 @@ where pub const trait Residual: 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; } diff --git a/library/coretests/tests/array.rs b/library/coretests/tests/array.rs index bdea7d4c0bdb..2b4429092e98 100644 --- a/library/coretests/tests/array.rs +++ b/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])); }