optimize zipping over array iterators

library/core/src/array/iter.rs

@@ -4,7 +4,7 @@ use crate::num::NonZeroUsize;
 use crate::{
  fmt,
  intrinsics::transmute_unchecked,
- iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
+ iter::{self, ExactSizeIterator, FusedIterator, TrustedLen, TrustedRandomAccessNoCoerce},
  mem::MaybeUninit,
  ops::{IndexRange, Range},
  ptr,
@@ -293,6 +293,12 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
 
  NonZeroUsize::new(remaining).map_or(Ok(()), Err)
  }
+
+ #[inline]
+ unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item {
+ // SAFETY: The caller must provide an idx that is in bound of the remainder.
+ unsafe { self.data.as_ptr().add(self.alive.start()).add(idx).cast::<T>().read() }
+ }
 }
 
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
@@ -374,6 +380,25 @@ impl<T, const N: usize> FusedIterator for IntoIter<T, N> {}
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
 unsafe impl<T, const N: usize> TrustedLen for IntoIter<T, N> {}
 
+#[doc(hidden)]
+#[unstable(issue = "none", feature = "std_internals")]
+#[rustc_unsafe_specialization_marker]
+pub trait NonDrop {}
+
+// T: Copy as approximation for !Drop since get_unchecked does not advance self.alive
+// and thus we can't implement drop-handling
+#[unstable(issue = "none", feature = "std_internals")]
+impl<T: Copy> NonDrop for T {}
+
+#[doc(hidden)]
+#[unstable(issue = "none", feature = "std_internals")]
+unsafe impl<T, const N: usize> TrustedRandomAccessNoCoerce for IntoIter<T, N>
+where
+ T: NonDrop,
+{
+ const MAY_HAVE_SIDE_EFFECT: bool = false;
+}
+
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
 impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
  fn clone(&self) -> Self {

library/core/src/iter/adapters/zip.rs

@@ -94,6 +94,14 @@ where
  ZipImpl::nth(self, n)
  }
 
+ #[inline]
+ fn fold<Acc, F>(self, init: Acc, f: F) -> Acc
+ where
+ F: FnMut(Acc, Self::Item) -> Acc,
+ {
+ ZipImpl::fold(self, init, f)
+ }
+
  #[inline]
  unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
  where
@@ -129,6 +137,9 @@ trait ZipImpl<A, B> {
  where
  A: DoubleEndedIterator + ExactSizeIterator,
  B: DoubleEndedIterator + ExactSizeIterator;
+ fn fold<Acc, F>(self, init: Acc, f: F) -> Acc
+ where
+ F: FnMut(Acc, Self::Item) -> Acc;
  // This has the same safety requirements as `Iterator::__iterator_get_unchecked`
  unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item
  where
@@ -228,6 +239,14 @@ where
  {
  unreachable!("Always specialized");
  }
+
+ #[inline]
+ default fn fold<Acc, F>(self, init: Acc, f: F) -> Acc
+ where
+ F: FnMut(Acc, Self::Item) -> Acc,
+ {
+ SpecFold::spec_fold(self, init, f)
+ }
 }
 
 #[doc(hidden)]
@@ -251,6 +270,24 @@ where
  // `Iterator::__iterator_get_unchecked`.
  unsafe { (self.a.__iterator_get_unchecked(idx), self.b.__iterator_get_unchecked(idx)) }
  }
+
+ #[inline]
+ fn fold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
+ where
+ F: FnMut(Acc, Self::Item) -> Acc,
+ {
+ let mut accum = init;
+ let len = ZipImpl::size_hint(&self).0;
+ for i in 0..len {
+ // SAFETY: since Self: TrustedRandomAccessNoCoerce we can trust the size-hint to
+ // calculate the length and then use that to do unchecked iteration.
+ // fold consumes the iterator so we don't need to fixup any state.
+ unsafe {
+ accum = f(accum, self.get_unchecked(i));
+ }
+ }
+ accum
+ }
 }
 
 #[doc(hidden)]
@@ -590,3 +627,56 @@ unsafe impl<I: Iterator + TrustedRandomAccessNoCoerce> SpecTrustedRandomAccess f
  unsafe { self.__iterator_get_unchecked(index) }
  }
 }
+
+trait SpecFold: Iterator {
+ fn spec_fold<B, F>(self, init: B, f: F) -> B
+ where
+ Self: Sized,
+ F: FnMut(B, Self::Item) -> B;
+}
+
+impl<A: Iterator, B: Iterator> SpecFold for Zip<A, B> {
+ // Adapted from default impl from the Iterator trait
+ #[inline]
+ default fn spec_fold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
+ where
+ F: FnMut(Acc, Self::Item) -> Acc,
+ {
+ let mut accum = init;
+ while let Some(x) = ZipImpl::next(&mut self) {
+ accum = f(accum, x);
+ }
+ accum
+ }
+}
+
+impl<A: TrustedLen, B: TrustedLen> SpecFold for Zip<A, B> {
+ #[inline]
+ fn spec_fold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
+ where
+ F: FnMut(Acc, Self::Item) -> Acc,
+ {
+ let mut accum = init;
+ loop {
+ let (upper, more) = if let Some(upper) = ZipImpl::size_hint(&self).1 {
+ (upper, false)
+ } else {
+ // Per TrustedLen contract a None upper bound means more than usize::MAX items
+ (usize::MAX, true)
+ };
+
+ for _ in 0..upper {
+ let pair =
+ // SAFETY: TrustedLen guarantees that at least `upper` many items are available
+ // therefore we know they can't be None
+ unsafe { (self.a.next().unwrap_unchecked(), self.b.next().unwrap_unchecked()) };
+ accum = f(accum, pair);
+ }
+
+ if !more {
+ break;
+ }
+ }
+ accum
+ }
+}

library/core/tests/iter/adapters/zip.rs

@@ -184,7 +184,11 @@ fn test_zip_nested_sideffectful() {
  let it = xs.iter_mut().map(|x| *x = 1).enumerate().zip(&ys);
  it.count();
  }
- assert_eq!(&xs, &[1, 1, 1, 1, 1, 0]);
+ let length_aware = &xs == &[1, 1, 1, 1, 0, 0];
+ let probe_first = &xs == &[1, 1, 1, 1, 1, 0];
+
+ // either implementation is valid according to zip documentation
+ assert!(length_aware || probe_first);
 }
 
 #[test]

tests/assembly/libs/issue-115339-zip-arrays.rs

@@ -0,0 +1,25 @@
+// assembly-output: emit-asm
+// # zen3 previously exhibited odd vectorization
+// compile-flags: --crate-type=lib -Ctarget-cpu=znver3 -O
+// only-x86_64
+// ignore-sgx
+
+use std::iter;
+
+// previously this produced a long chain of
+// 56: vpextrb $6, %xmm0, %ecx
+// 57: orb %cl, 22(%rsi)
+// 58: vpextrb $7, %xmm0, %ecx
+// 59: orb %cl, 23(%rsi)
+// [...]
+
+// CHECK-LABEL: zip_arrays:
+#[no_mangle]
+pub fn zip_arrays(mut a: [u8; 32], b: [u8; 32]) -> [u8; 32] {
+ // CHECK-NOT: vpextrb
+ // CHECK-NOT: orb %cl
+ // CHECK: vorps
+ iter::zip(&mut a, b).for_each(|(a, b)| *a |= b);
+ // CHECK: retq
+ a
+}