array zip_map feature

library/core/src/array/iter.rs

@@ -3,7 +3,7 @@
 use crate::{
     cmp, fmt,
     iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
-    mem::{self, MaybeUninit},
+    mem::{self, ManuallyDrop, MaybeUninit},
     ops::Range,
     ptr,
 };
@@ -233,6 +233,39 @@ impl<T, const N: usize> IntoIter<T, N> {
             MaybeUninit::slice_assume_init_mut(slice)
         }
     }
+
+    pub(super) unsafe fn pop_front_unchecked(&mut self) -> T {
+        debug_assert!(!self.alive.is_empty());
+        debug_assert!(self.alive.start < N);
+
+        // SAFETY: The caller should ensure that these values are still valid
+        unsafe {
+            let front = MaybeUninit::assume_init_read(self.data.get_unchecked(self.alive.start));
+            self.alive.start += 1;
+            front
+        }
+    }
+
+    pub(super) unsafe fn push_unchecked(&mut self, value: T) {
+        debug_assert!(self.alive.end < N);
+
+        // SAFETY: The caller should ensure that there is still free space
+        unsafe {
+            self.data.get_unchecked_mut(self.alive.end).write(value);
+            self.alive.end += 1;
+        }
+    }
+
+    pub(super) unsafe fn assume_init(self) -> [T; N] {
+        debug_assert!(self.alive == (0..N));
+        let this = ManuallyDrop::new(self);
+
+        // SAFETY: The caller should ensure that the array is full
+        unsafe {
+            let data = core::ptr::read(&this.data);
+            MaybeUninit::array_assume_init(data)
+        }
+    }
 }
 
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]

library/core/src/array/mod.rs

@@ -503,11 +503,121 @@ impl<T, const N: usize> [T; N] {
     /// ```
     #[unstable(feature = "array_zip", issue = "80094")]
     pub fn zip<U>(self, rhs: [U; N]) -> [(T, U); N] {
-        let mut iter = IntoIterator::into_iter(self).zip(rhs);
+        self.zip_map(rhs, |a, b| (a, b))
+    }
 
-        // SAFETY: we know for certain that this iterator will yield exactly `N`
-        // items.
-        unsafe { collect_into_array_unchecked(&mut iter) }
+    /// 'Zips up' two arrays into a single array, applying the `op` over the pairs
+    ///
+    /// This is equivalent but faster than doing manual zip + map
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// #![feature(array_zip, array_zip_map)]
+    /// # use std::ops::Add;
+    /// let op = i32::add;
+    ///
+    /// let x = [1, 2, 3];
+    /// let y = [4, 5, 6];
+    ///
+    /// let output1 = x.zip(y).map(|(a, b)| op(a, b));
+    /// let output2 = x.zip_map(y, op);
+    ///
+    /// assert_eq!(output1, output2);
+    /// ```
+    #[unstable(feature = "array_zip_map", issue = "none")]
+    pub fn zip_map<U, F, R>(self, rhs: [U; N], mut op: F) -> [R; N]
+    where
+        F: FnMut(T, U) -> R,
+    {
+        // If any of the items need drop, take the 'slow' path that ensures all drop
+        // handlers are called in case of panic.
+        //
+        // In my profiling, I found that the drop code can interfere with the codegen just enough
+        // to cause some noticable regressions in speed
+        if core::mem::needs_drop::<T>()
+            || core::mem::needs_drop::<U>()
+            || core::mem::needs_drop::<R>()
+        {
+            let mut lhs = IntoIterator::into_iter(self);
+            let mut rhs = IntoIterator::into_iter(rhs);
+            let mut output = IntoIter::empty();
+
+            for _ in 0..N {
+                // SAFETY: Will only be called a maximum of N times
+                unsafe {
+                    let lhs = lhs.pop_front_unchecked();
+                    let rhs = rhs.pop_front_unchecked();
+                    output.push_unchecked(op(lhs, rhs));
+                }
+            }
+
+            // SAFETY: We have confirmed it has N elements
+            unsafe { output.assume_init() }
+        } else {
+            let mut output: [MaybeUninit<R>; N] = MaybeUninit::uninit_array();
+
+            for i in 0..N {
+                // SAFETY: Will only be called a maximum of N times
+                unsafe {
+                    let lhs = core::ptr::read(&self[i]);
+                    let rhs = core::ptr::read(&rhs[i]);
+                    output[i].write(op(lhs, rhs));
+                }
+            }
+
+            // SAFETY: We have confirmed it has N elements
+            unsafe { MaybeUninit::array_assume_init(output) }
+        }
+    }
+
+    /// Applying the `op` over pairs of elements and assigning the value
+    /// back into `self`
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// #![feature(array_zip_map)]
+    /// # use std::ops::AddAssign;
+    ///
+    /// let mut x = [1, 2, 3];
+    /// let y = [4, 5, 6];
+    ///
+    /// x.zip_map_assign(y, i32::add_assign);
+    ///
+    /// assert_eq!(x, [5, 7, 9]);
+    /// ```
+    #[unstable(feature = "array_zip_map", issue = "none")]
+    pub fn zip_map_assign<U, F>(&mut self, rhs: [U; N], mut op: F)
+    where
+        F: FnMut(&mut T, U),
+    {
+        // If the rhs need drop, take the 'slow' path that ensures all drop
+        // handlers are called in case of panic.
+        //
+        // In my profiling, I found that the drop code can interfere with the codegen just enough
+        // to cause some noticable regressions in speed
+        if core::mem::needs_drop::<U>() {
+            let mut rhs = IntoIterator::into_iter(rhs);
+
+            for i in 0..N {
+                // SAFETY: Will only be called a maximum of N times
+                unsafe {
+                    let lhs = self.get_unchecked_mut(i);
+                    let rhs = rhs.pop_front_unchecked();
+                    op(lhs, rhs)
+                }
+            }
+        } else {
+            for i in 0..N {
+                // SAFETY: Will only be called a maximum of N times
+                unsafe {
+                    let lhs = self.get_unchecked_mut(i);
+                    let rhs = core::ptr::read(&rhs[i]);
+                    op(lhs, rhs)
+                }
+            }
+        }
     }
 
     /// Returns a slice containing the entire array. Equivalent to `&s[..]`.

library/core/tests/array.rs

@@ -1,5 +1,7 @@
 use core::array;
 use core::convert::TryFrom;
+use core::iter::Iterator;
+use core::ops::{Add, AddAssign};
 use core::sync::atomic::{AtomicUsize, Ordering};
 
 #[test]
@@ -668,3 +670,43 @@ fn array_mixed_equality_nans() {
     assert!(!(mut3 == array3));
     assert!(mut3 != array3);
 }
+
+#[test]
+fn add() {
+    let a: [i32; 4] = [0, 1, 2, 3];
+    let b: [i32; 4] = [4, 5, 6, 7];
+
+    assert_eq!(a.zip_map(b, Add::add), [4, 6, 8, 10]);
+}
+
+#[test]
+fn add_assign() {
+    let mut a: [i32; 4] = [0, 1, 2, 3];
+    let b: [i32; 4] = [4, 5, 6, 7];
+
+    a.zip_map_assign(b, AddAssign::add_assign);
+
+    assert_eq!(a, [4, 6, 8, 10]);
+}
+
+#[test]
+fn add_many() {
+    let a: [i32; 128] = (0..128).collect::<Vec<_>>().try_into().unwrap();
+    let b: [i32; 128] = (128..256).collect::<Vec<_>>().try_into().unwrap();
+
+    let exp: [i32; 128] = (128..384).step_by(2).collect::<Vec<_>>().try_into().unwrap();
+
+    assert_eq!(a.zip_map(b, Add::add), exp);
+}
+
+#[test]
+fn add_assign_many() {
+    let mut a: [i32; 128] = (0..128).collect::<Vec<_>>().try_into().unwrap();
+    let b: [i32; 128] = (128..256).collect::<Vec<_>>().try_into().unwrap();
+
+    let exp: [i32; 128] = (128..384).step_by(2).collect::<Vec<_>>().try_into().unwrap();
+
+    a.zip_map_assign(b, AddAssign::add_assign);
+
+    assert_eq!(a, exp);
+}

library/core/tests/lib.rs

@@ -2,6 +2,7 @@
 #![feature(array_chunks)]
 #![feature(array_methods)]
 #![feature(array_windows)]
+#![feature(array_zip_map)]
 #![feature(bench_black_box)]
 #![feature(bool_to_option)]
 #![feature(box_syntax)]