Add in-place optimization for array map

Author: JulianKnodt

library/core/src/array/mod.rs

@@ -429,7 +429,68 @@ impl<T, const N: usize> [T; N] {
     where
         F: FnMut(T) -> U,
     {
-        use crate::mem::MaybeUninit;
+        use crate::mem::{align_of, forget, size_of, transmute_copy, ManuallyDrop, MaybeUninit};
+
+        if align_of::<T>() == align_of::<U>() && size_of::<T>() == size_of::<U>() {
+            // this branch allows reuse of the original array as a backing store
+            // kind of. As written with no compiler optimizations, transmute copy will
+            // still require 2 copies of the original array, but when it can be converted to
+            // transmute, this will require 0 copies.
+            union Translated<T, U> {
+                src: MaybeUninit<T>,
+                dst: ManuallyDrop<U>,
+            };
+            struct Guard<T, U, const N: usize> {
+                data: *mut [Translated<T, U>; N],
+                initialized: usize,
+            }
+            impl<T, U, const N: usize> Drop for Guard<T, U, N> {
+                fn drop(&mut self) {
+                    debug_assert!(self.initialized < N);
+                    let initialized_part =
+                        crate::ptr::slice_from_raw_parts_mut(self.data as *mut U, self.initialized);
+                    // SAFETY:
+                    // since we read from the element at initialized then panicked,
+                    // we have to skip over it to not double drop.
+                    let todo_ptr = unsafe { self.data.add(self.initialized + 1) as *mut T };
+                    let todo_part =
+                        crate::ptr::slice_from_raw_parts_mut(todo_ptr, N - self.initialized - 1);
+                    // SAFETY:
+                    // Have to remove both the initialized and not yet reached items.
+                    unsafe {
+                        crate::ptr::drop_in_place(initialized_part);
+                        crate::ptr::drop_in_place(todo_part);
+                    }
+                }
+            }
+            // SAFETY:
+            // Since we know that T & U have the same size and alignment we can safely transmute
+            // between them here
+            let mut src_dst = unsafe { transmute_copy::<_, [Translated<T, U>; N]>(&self) };
+
+            let mut guard: Guard<T, U, N> = Guard { data: &mut src_dst, initialized: 0 };
+            // Need to forget self now because the guard is responsible for dropping the items
+            forget(self);
+            for i in 0..N {
+                // SAFETY:
+                // All items prior to `i` are the `dst` variant.
+                // In order to convert `i` from src to dst, we take it from `MaybeUninit`,
+                // leaving uninitialized in its place, and set the destination as
+                // ManuallyDrop::new(..), and implicitly know that it will be a `dst` variant
+                // from where
+                unsafe {
+                    let v = f(src_dst[i].src.read());
+                    src_dst[i].dst = ManuallyDrop::new(v);
+                }
+                guard.initialized += 1;
+            }
+            forget(guard);
+            // SAFETY:
+            // At this point all the items have been initialized and are in `dst` discriminant.
+            // We can switch them over to being of type `U`.
+            return unsafe { transmute_copy::<_, [U; N]>(&src_dst) };
+        }
+
         struct Guard<T, const N: usize> {
             dst: *mut T,
             initialized: usize,
@@ -457,10 +518,10 @@ impl<T, const N: usize> [T; N] {
         }
         // FIXME: Convert to crate::mem::transmute once it works with generics.
         // unsafe { crate::mem::transmute::<[MaybeUninit<U>; N], [U; N]>(dst) }
-        crate::mem::forget(guard);
+        forget(guard);
         // SAFETY: At this point we've properly initialized the whole array
         // and we just need to cast it to the correct type.
-        unsafe { crate::mem::transmute_copy::<_, [U; N]>(&dst) }
+        unsafe { transmute_copy::<_, [U; N]>(&dst) }
     }
 
     /// Returns a slice containing the entire array. Equivalent to `&s[..]`.

library/core/tests/array.rs

@@ -323,26 +323,53 @@ fn array_map() {
 fn array_map_drop_safety() {
     use core::sync::atomic::AtomicUsize;
     use core::sync::atomic::Ordering;
-    static DROPPED: AtomicUsize = AtomicUsize::new(0);
-    struct DropCounter;
-    impl Drop for DropCounter {
+    static DROPPED: [AtomicUsize; 3] =
+        [AtomicUsize::new(0), AtomicUsize::new(0), AtomicUsize::new(0)];
+    struct DropCounter<const N: usize>;
+    impl<const N: usize> Drop for DropCounter<N> {
         fn drop(&mut self) {
-            DROPPED.fetch_add(1, Ordering::SeqCst);
+            DROPPED[N].fetch_add(1, Ordering::SeqCst);
         }
     }
 
-    let num_to_create = 5;
-    let success = std::panic::catch_unwind(|| {
-        let items = [0; 10];
-        let mut nth = 0;
-        items.map(|_| {
-            assert!(nth < num_to_create);
-            nth += 1;
-            DropCounter
+    {
+        let num_to_create = 5;
+        let success = std::panic::catch_unwind(|| {
+            let items = [0; 10];
+            let mut nth = 0;
+            items.map(|_| {
+                assert!(nth < num_to_create);
+                nth += 1;
+                DropCounter::<0>
+            });
+        });
+        assert!(success.is_err());
+        assert_eq!(DROPPED[0].load(Ordering::SeqCst), num_to_create);
+    }
+
+    {
+        assert_eq!(DROPPED[1].load(Ordering::SeqCst), 0);
+        let num_to_create = 3;
+        const TOTAL: usize = 5;
+        let success = std::panic::catch_unwind(|| {
+            let items: [DropCounter<1>; TOTAL] = [
+                DropCounter::<1>,
+                DropCounter::<1>,
+                DropCounter::<1>,
+                DropCounter::<1>,
+                DropCounter::<1>,
+            ];
+            let mut nth = 0;
+            items.map(|_| {
+                assert!(nth < num_to_create);
+                nth += 1;
+                DropCounter::<2>
+            });
         });
-    });
-    assert!(success.is_err());
-    assert_eq!(DROPPED.load(Ordering::SeqCst), num_to_create);
+        assert!(success.is_err());
+        assert_eq!(DROPPED[2].load(Ordering::SeqCst), num_to_create);
+        assert_eq!(DROPPED[1].load(Ordering::SeqCst), TOTAL);
+    }
     panic!("test succeeded")
 }
 

library/core/tests/lib.rs

@@ -32,6 +32,7 @@
 #![feature(raw)]
 #![feature(sort_internals)]
 #![feature(slice_partition_at_index)]
+#![feature(min_const_generics)]
 #![feature(min_specialization)]
 #![feature(step_trait)]
 #![feature(step_trait_ext)]