From 963d6f757c52dd0693e3c0541d7b027279f8a2db Mon Sep 17 00:00:00 2001 From: The 8472 Date: Mon, 17 Oct 2022 22:47:39 +0200 Subject: [PATCH 1/2] add a benchmark for slice_iter.copied().array_chunks() --- library/core/benches/iter.rs | 20 ++++++++++++++++++++ library/core/benches/lib.rs | 1 + 2 files changed, 21 insertions(+) diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index 4b40485d207c8..38887f29af153 100644 --- a/library/core/benches/iter.rs +++ b/library/core/benches/iter.rs @@ -1,4 +1,6 @@ use core::iter::*; +use core::mem; +use core::num::Wrapping; use test::{black_box, Bencher}; #[bench] @@ -398,3 +400,21 @@ fn bench_trusted_random_access_adapters(b: &mut Bencher) { acc }) } + +/// Exercises the iter::Copied specialization for slice::Iter +#[bench] +fn bench_copied_array_chunks(b: &mut Bencher) { + let v = vec![1u8; 1024]; + + b.iter(|| { + black_box(&v) + .iter() + .copied() + .array_chunks::<{ mem::size_of::() }>() + .map(|ary| { + let d = u64::from_ne_bytes(ary); + Wrapping(d.rotate_left(7).wrapping_add(1)) + }) + .sum::>() + }) +} diff --git a/library/core/benches/lib.rs b/library/core/benches/lib.rs index a6c174d2fca21..1e462e3fc3f8c 100644 --- a/library/core/benches/lib.rs +++ b/library/core/benches/lib.rs @@ -4,6 +4,7 @@ #![feature(int_log)] #![feature(test)] #![feature(trusted_random_access)] +#![feature(iter_array_chunks)] extern crate test; From 873a18e2213a506d4cf6357e5faef0990ec08683 Mon Sep 17 00:00:00 2001 From: The 8472 Date: Mon, 17 Oct 2022 22:49:32 +0200 Subject: [PATCH 2/2] specialize slice_iter.copied().next_chunk() --- library/core/src/iter/adapters/copied.rs | 74 ++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/library/core/src/iter/adapters/copied.rs b/library/core/src/iter/adapters/copied.rs index f9bfd77d7fbbf..62d3afb81603d 100644 --- a/library/core/src/iter/adapters/copied.rs +++ b/library/core/src/iter/adapters/copied.rs @@ -2,7 +2,10 @@ use crate::iter::adapters::{ zip::try_get_unchecked, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; use crate::iter::{FusedIterator, TrustedLen}; +use crate::mem::MaybeUninit; +use crate::mem::SizedTypeProperties; use crate::ops::Try; +use crate::{array, ptr}; /// An iterator that copies the elements of an underlying iterator. /// @@ -44,6 +47,15 @@ where self.it.next().copied() } + fn next_chunk( + &mut self, + ) -> Result<[Self::Item; N], array::IntoIter> + where + Self: Sized, + { + >::spec_next_chunk(&mut self.it) + } + fn size_hint(&self) -> (usize, Option) { self.it.size_hint() } @@ -166,3 +178,65 @@ where T: Copy, { } + +trait SpecNextChunk<'a, const N: usize, T: 'a>: Iterator +where + T: Copy, +{ + fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter>; +} + +impl<'a, const N: usize, I, T: 'a> SpecNextChunk<'a, N, T> for I +where + I: Iterator, + T: Copy, +{ + default fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter> { + array::iter_next_chunk(&mut self.map(|e| *e)) + } +} + +impl<'a, const N: usize, T: 'a> SpecNextChunk<'a, N, T> for crate::slice::Iter<'a, T> +where + T: Copy, +{ + fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter> { + let mut raw_array = MaybeUninit::uninit_array(); + + let len = self.len(); + + if T::IS_ZST { + if len < N { + let _ = self.advance_by(len); + // SAFETY: ZSTs can be conjured ex nihilo; only the amount has to be correct + return Err(unsafe { array::IntoIter::new_unchecked(raw_array, 0..len) }); + } + + let _ = self.advance_by(N); + // SAFETY: ditto + return Ok(unsafe { MaybeUninit::array_assume_init(raw_array) }); + } + + if len < N { + // SAFETY: `len` indicates that this many elements are available and we just checked that + // it fits into the array. + unsafe { + ptr::copy_nonoverlapping( + self.as_ref().as_ptr(), + raw_array.as_mut_ptr() as *mut T, + len, + ); + let _ = self.advance_by(len); + return Err(array::IntoIter::new_unchecked(raw_array, 0..len)); + } + } + + // SAFETY: `len` is larger than the array size. Copy a fixed amount here to fully initialize + // the array. + unsafe { + ptr::copy_nonoverlapping(self.as_ref().as_ptr(), raw_array.as_mut_ptr() as *mut T, N); + let _ = self.advance_by(N); + Ok(MaybeUninit::array_assume_init(raw_array)) + } + } +}