core: drop unmapped ZSTs in array map

Author: joboet

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,44 +14,47 @@ 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<const N: usize>(
+        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 }
+            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.
 #[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.
     ///
     /// 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,
+    /// 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<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, '_, T, N, F>
+impl<T, U, F> 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<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, '_, T, N, F>
+impl<T, U, F> 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::<T>() })
@@ -89,20 +102,32 @@ where
 }
 #[rustc_const_unstable(feature = "array_try_map", issue = "79711")]
 #[unstable(feature = "array_try_map", issue = "79711")]
-impl<T: [const] Destruct, const N: usize, F> const Drop for Drain<'_, '_, T, N, F> {
+impl<T: [const] Destruct, F> 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) }
     }
 }

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<Cell<usize>> = 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() {

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)]