diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs index e5024c215be9c..c02083a3f85e4 100644 --- a/library/core/src/array/iter.rs +++ b/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 IntoIter { 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")] diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index ee79021ed536e..607a12be5e663 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -503,11 +503,121 @@ impl [T; N] { /// ``` #[unstable(feature = "array_zip", issue = "80094")] pub fn zip(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(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::() + || core::mem::needs_drop::() + || core::mem::needs_drop::() + { + 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; 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(&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::() { + 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[..]`. diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs index a778779c0fd88..f60d28a9a4a5f 100644 --- a/library/core/tests/array.rs +++ b/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::>().try_into().unwrap(); + let b: [i32; 128] = (128..256).collect::>().try_into().unwrap(); + + let exp: [i32; 128] = (128..384).step_by(2).collect::>().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::>().try_into().unwrap(); + let b: [i32; 128] = (128..256).collect::>().try_into().unwrap(); + + let exp: [i32; 128] = (128..384).step_by(2).collect::>().try_into().unwrap(); + + a.zip_map_assign(b, AddAssign::add_assign); + + assert_eq!(a, exp); +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 28830222c1a94..b15078634def5 100644 --- a/library/core/tests/lib.rs +++ b/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)]