diff --git a/library/core/src/array/drain.rs b/library/core/src/array/drain.rs index 17792dca583d2..b2ff54bdfa21c 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,16 +26,17 @@ 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 } } } } /// See [`Drain::new`]; this is our fake iterator. #[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. /// @@ -41,16 +44,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, { @@ -63,7 +66,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, { @@ -73,6 +76,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::() }) @@ -88,20 +101,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 43fed944e9280..a3b0e59278f79 100644 --- a/library/coretests/tests/array.rs +++ b/library/coretests/tests/array.rs @@ -1,3 +1,4 @@ +use core::cell::Cell; use core::num::NonZero; use core::sync::atomic::{AtomicUsize, Ordering}; use core::{array, assert_eq}; @@ -168,8 +169,6 @@ fn iterator_debug() { #[test] fn iterator_drops() { - use core::cell::Cell; - // This test makes sure the correct number of elements are dropped. The `R` // type is just a reference to a `Cell` that is incremented when an `R` is // dropped. @@ -337,8 +336,6 @@ fn array_map_drop_safety() { #[test] fn cell_allows_array_cycle() { - use core::cell::Cell; - #[derive(Debug)] struct B<'a> { a: [Cell>>; 2], @@ -513,7 +510,6 @@ fn array_rsplit_array_mut_out_of_bounds() { #[test] fn array_intoiter_advance_by() { - use std::cell::Cell; struct DropCounter<'a>(usize, &'a Cell); impl Drop for DropCounter<'_> { fn drop(&mut self) { @@ -566,7 +562,6 @@ fn array_intoiter_advance_by() { #[test] fn array_intoiter_advance_back_by() { - use std::cell::Cell; struct DropCounter<'a>(usize, &'a Cell); impl Drop for DropCounter<'_> { fn drop(&mut self) { @@ -718,6 +713,33 @@ fn array_map_drops_unmapped_elements_on_panic() { } } +#[cfg(not(panic = "abort"))] +#[test] +fn array_map_drops_unmapped_zst_elements_on_panic() { + use std::sync::ReentrantLock; + + 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 c4292c2a421b1..47a21397a20f2 100644 --- a/library/coretests/tests/lib.rs +++ b/library/coretests/tests/lib.rs @@ -96,6 +96,7 @@ #![feature(pointer_is_aligned_to)] #![feature(portable_simd)] #![feature(ptr_metadata)] +#![feature(reentrant_lock)] #![feature(result_option_map_or_default)] #![feature(rustc_attrs)] #![feature(signed_bigint_helpers)]