Auto merge of rust-lang#75212 - JulianKnodt:array_map, r=LukasKalbertodt

Add `array` lang item and `[T; N]::map(f: FnMut(T) -> S)`

This introduces an `array` lang item so functions can be defined on top of `[T; N]`. This was previously not done because const-generics was not complete enough to allow for this. Now it is in a state that is usable enough to start adding functions.

The function added is a monadic (I think?) map from `[T; N] -> [S; N]`. Until transmute can function on arrays, it also allocates an extra temporary array, but this can be removed at some point.

r? @lcnr

Author: bors

core/src/array/mod.rs

@@ -364,3 +364,66 @@ macro_rules! array_impl_default {
 }
 
 array_impl_default! {32, T T T T T T T T T T T T T T T T T T T T T T T T T T T T T T T T}
+
+#[cfg(not(bootstrap))]
+#[lang = "array"]
+impl<T, const N: usize> [T; N] {
+    /// Returns an array of the same size as `self`, with function `f` applied to each element
+    /// in order.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(array_map)]
+    /// let x = [1, 2, 3];
+    /// let y = x.map(|v| v + 1);
+    /// assert_eq!(y, [2, 3, 4]);
+    ///
+    /// let x = [1, 2, 3];
+    /// let mut temp = 0;
+    /// let y = x.map(|v| { temp += 1; v * temp });
+    /// assert_eq!(y, [1, 4, 9]);
+    ///
+    /// let x = ["Ferris", "Bueller's", "Day", "Off"];
+    /// let y = x.map(|v| v.len());
+    /// assert_eq!(y, [6, 9, 3, 3]);
+    /// ```
+    #[unstable(feature = "array_map", issue = "75243")]
+    pub fn map<F, U>(self, mut f: F) -> [U; N]
+    where
+        F: FnMut(T) -> U,
+    {
+        use crate::mem::MaybeUninit;
+        struct Guard<T, const N: usize> {
+            dst: *mut T,
+            initialized: usize,
+        }
+
+        impl<T, const N: usize> Drop for Guard<T, N> {
+            fn drop(&mut self) {
+                debug_assert!(self.initialized <= N);
+
+                let initialized_part =
+                    crate::ptr::slice_from_raw_parts_mut(self.dst, self.initialized);
+                // SAFETY: this raw slice will contain only initialized objects
+                // that's why, it is allowed to drop it.
+                unsafe {
+                    crate::ptr::drop_in_place(initialized_part);
+                }
+            }
+        }
+        let mut dst = MaybeUninit::uninit_array::<N>();
+        let mut guard: Guard<U, N> =
+            Guard { dst: MaybeUninit::first_ptr_mut(&mut dst), initialized: 0 };
+        for (src, dst) in IntoIter::new(self).zip(&mut dst) {
+            dst.write(f(src));
+            guard.initialized += 1;
+        }
+        // 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);
+        // 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) }
+    }
+}

core/src/lib.rs

@@ -145,6 +145,7 @@
 #![feature(abi_unadjusted)]
 #![feature(adx_target_feature)]
 #![feature(maybe_uninit_slice)]
+#![feature(maybe_uninit_extra)]
 #![feature(external_doc)]
 #![feature(associated_type_bounds)]
 #![feature(const_caller_location)]

core/tests/array.rs

@@ -290,3 +290,43 @@ fn empty_array_is_always_default() {
 
     let _arr = <[DoesNotImplDefault; 0]>::default();
 }
+
+#[test]
+fn array_map() {
+    let a = [1, 2, 3];
+    let b = a.map(|v| v + 1);
+    assert_eq!(b, [2, 3, 4]);
+
+    let a = [1u8, 2, 3];
+    let b = a.map(|v| v as u64);
+    assert_eq!(b, [1, 2, 3]);
+}
+
+// See note on above test for why `should_panic` is used.
+#[test]
+#[should_panic(expected = "test succeeded")]
+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 {
+        fn drop(&mut self) {
+            DROPPED.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
+        });
+    });
+    assert!(success.is_err());
+    assert_eq!(DROPPED.load(Ordering::SeqCst), num_to_create);
+    panic!("test succeeded")
+}

core/tests/lib.rs

@@ -1,5 +1,6 @@
 #![feature(alloc_layout_extra)]
 #![feature(array_chunks)]
+#![feature(array_map)]
 #![feature(bool_to_option)]
 #![feature(bound_cloned)]
 #![feature(box_syntax)]