diff --git a/library/core/src/array/drain.rs b/library/core/src/array/drain.rs index 1c6137191324c..54fc080e24528 100644 --- a/library/core/src/array/drain.rs +++ b/library/core/src/array/drain.rs @@ -1,8 +1,8 @@ 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}; +use crate::mem::{ManuallyDrop, SizedTypeProperties, conjure_zst, transmute}; +use crate::ptr::{NonNull, drop_in_place, from_raw_parts_mut, without_provenance_mut}; -impl<'l, 'f, T, U, const N: usize, F: FnMut(T) -> U> Drain<'l, 'f, T, N, F> { +impl<'l, 'f, T, U, F: FnMut(T) -> U> Drain<'l, 'f, T, 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 @@ -14,9 +14,11 @@ impl<'l, 'f, T, U, const N: usize, F: FnMut(T) -> U> Drain<'l, 'f, T, N, F> { /// 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 { + 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: @@ -24,8 +26,9 @@ impl<'l, 'f, T, U, const N: usize, F: FnMut(T) -> U> Drain<'l, 'f, T, N, F> { // 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 } + let end_or_len = + if T::IS_ZST { without_provenance_mut(N) } else { ptr.as_ptr().add(N) }; + Self { ptr, end_or_len, f, l: PhantomData } } } } @@ -33,8 +36,8 @@ impl<'l, 'f, T, U, const N: usize, F: FnMut(T) -> U> Drain<'l, 'f, T, N, F> { /// 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. +pub(super) struct Drain<'l, 'f, T, F> { + // FIXME(const-hack): This is a slice::IterMut<'l>, replace when possible. /// The pointer to the next element to return, or the past-the-end location /// if the drainer is empty. /// @@ -42,16 +45,16 @@ pub(super) struct Drain<'l, 'f, T, const N: usize, F> { /// 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, + /// For ZSTs, this is the number of unprocessed items. + end_or_len: *mut T, f: &'f mut F, - l: PhantomData<&'l mut [T; N]>, + l: PhantomData<&'l mut [T]>, } #[rustc_const_unstable(feature = "array_try_map", issue = "79711")] #[unstable(feature = "array_try_map", issue = "79711")] -impl const FnOnce<(usize,)> for &mut Drain<'_, '_, T, N, F> +impl const FnOnce<(usize,)> for &mut Drain<'_, '_, T, F> where F: [const] FnMut(T) -> U, { @@ -64,7 +67,7 @@ where } #[rustc_const_unstable(feature = "array_try_map", issue = "79711")] #[unstable(feature = "array_try_map", issue = "79711")] -impl const FnMut<(usize,)> for &mut Drain<'_, '_, T, N, F> +impl const FnMut<(usize,)> for &mut Drain<'_, '_, T, F> where F: [const] FnMut(T) -> U, { @@ -74,6 +77,16 @@ where (_ /* ignore argument */,): (usize,), ) -> Self::Output { if T::IS_ZST { + #[expect(ptr_to_integer_transmute_in_consts)] + // SAFETY: + // This is equivalent to `self.end_or_len.addr`, but that's not + // available in `const`. `self.end_or_len` doesn't have provenance, + // so transmuting is fine. + let len = unsafe { transmute::<*mut T, usize>(self.end_or_len) }; + // SAFETY: + // The caller guarantees that this is never called more than N times + // (see `Drain::new`), hence this cannot underflow. + self.end_or_len = without_provenance_mut(unsafe { len.unchecked_sub(1) }); // 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::() }) @@ -89,20 +102,32 @@ where } #[rustc_const_unstable(feature = "array_try_map", issue = "79711")] #[unstable(feature = "array_try_map", issue = "79711")] -impl const Drop for Drain<'_, '_, T, N, F> { +impl const Drop for Drain<'_, '_, T, F> { fn drop(&mut self) { - if !T::IS_ZST { + let slice = if T::IS_ZST { + from_raw_parts_mut::<[T]>( + self.ptr.as_ptr(), + #[expect(ptr_to_integer_transmute_in_consts)] + // SAFETY: + // This is equivalent to `self.end_or_len.addr`, but that's not + // available in `const`. `self.end_or_len` doesn't have provenance, + // so transmuting is fine. + unsafe { + transmute::<*mut T, usize>(self.end_or_len) + }, + ) + } else { // SAFETY: we cant read more than N elements - let slice = unsafe { + unsafe { from_raw_parts_mut::<[T]>( self.ptr.as_ptr(), // SAFETY: `start <= end` - self.end.offset_from_unsigned(self.ptr.as_ptr()), + self.end_or_len.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) } - } + // SAFETY: By the type invariant, we're allowed to drop all these. (we own it, after all) + unsafe { drop_in_place(slice) } } } diff --git a/library/coretests/tests/array.rs b/library/coretests/tests/array.rs index 2b4429092e98b..f68204de700a0 100644 --- a/library/coretests/tests/array.rs +++ b/library/coretests/tests/array.rs @@ -1,6 +1,8 @@ +use core::cell::Cell; use core::num::NonZero; use core::sync::atomic::{AtomicUsize, Ordering}; use core::{array, assert_eq}; +use std::sync::ReentrantLock; #[test] fn array_from_ref() { @@ -718,6 +720,31 @@ fn array_map_drops_unmapped_elements_on_panic() { } } +#[cfg(not(panic = "abort"))] +#[test] +fn array_map_drops_unmapped_zst_elements_on_panic() { + static DROPPED: ReentrantLock> = ReentrantLock::new(Cell::new(0)); + + struct ZstDrop; + impl Drop for ZstDrop { + fn drop(&mut self) { + DROPPED.lock().update(|x| x + 1); + } + } + + let dropped = DROPPED.lock(); + dropped.set(0); + let array = [const { ZstDrop }; 5]; + let success = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _ = array.map(|x| { + drop(x); + assert_eq!(dropped.get(), 1); + }); + })); + assert!(success.is_err()); + assert_eq!(dropped.get(), 5); +} + // This covers the `PartialEq::<[T]>::eq` impl for `[T; N]` when it returns false. #[test] fn array_eq() { diff --git a/library/coretests/tests/lib.rs b/library/coretests/tests/lib.rs index 5923328655524..bac4f3b228007 100644 --- a/library/coretests/tests/lib.rs +++ b/library/coretests/tests/lib.rs @@ -99,6 +99,7 @@ #![feature(pointer_is_aligned_to)] #![feature(portable_simd)] #![feature(ptr_metadata)] +#![feature(reentrant_lock)] #![feature(result_option_map_or_default)] #![feature(signed_bigint_helpers)] #![feature(slice_from_ptr_range)]