Skip to content

Commit

Permalink
Use next_chunk in ArrayChunks impl
Browse files Browse the repository at this point in the history
  • Loading branch information
@WaffleLapkin
WaffleLapkin committed Aug 2, 2022
1 parent de097a3 commit 69fc19a
Showing 1 changed file with 37 additions and 132 deletions.
169 changes: 37 additions & 132 deletions core/src/iter/adapters/array_chunks.rs
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,6 @@
use crate::array;
use crate::iter::{FusedIterator, Iterator};
use crate::mem;
use crate::mem::MaybeUninit;
use crate::ops::{ControlFlow, NeverShortCircuit, Try};
use crate::ptr;

/// An iterator over `N` elements of the iterator at a time.
///
Expand Down Expand Up @@ -70,37 +67,18 @@ where
F: FnMut(B, Self::Item) -> R,
R: Try<Output = B>,
{
let mut array = MaybeUninit::uninit_array();
// SAFETY: `array` will still be valid if `guard` is dropped.
let mut guard = unsafe { FrontGuard::new(&mut array) };

let result = self.iter.try_fold(init, |mut acc, item| {
// SAFETY: `init` starts at 0, increases by one each iteration and
// is reset to 0 once it reaches N.
unsafe { array.get_unchecked_mut(guard.init) }.write(item);
guard.init += 1;
if guard.init == N {
guard.init = 0;
let array = mem::replace(&mut array, MaybeUninit::uninit_array());
// SAFETY: the condition above asserts that all elements are
// initialized.
let item = unsafe { MaybeUninit::array_assume_init(array) };
acc = f(acc, item)?;
}
R::from_output(acc)
});
match result.branch() {
ControlFlow::Continue(o) => {
if guard.init > 0 {
let init = guard.init;
mem::forget(guard);
// SAFETY: `array` was initialized with `init` elements.
self.remainder =
Some(unsafe { array::IntoIter::new_unchecked(array, 0..init) });
let mut acc = init;
loop {
match self.iter.next_chunk() {
Ok(chunk) => acc = f(acc, chunk)?,
Err(remainder) => {
// Make sure to not override `self.remainder` with an empty array
// when `next` is called after `ArrayChunks` exhaustion.
self.remainder.get_or_insert(remainder);

break try { acc };
}
R::from_output(o)
}
ControlFlow::Break(r) => R::from_residual(r),
}
}

Expand All @@ -113,33 +91,6 @@ where
}
}

/// A guard for an array where elements are filled from the left.
struct FrontGuard<T, const N: usize> {
/// A pointer to the array that is being filled. We need to use a raw
/// pointer here because of the lifetime issues in the fold implementations.
ptr: *mut T,
/// The number of *initialized* elements.
init: usize,
}

impl<T, const N: usize> FrontGuard<T, N> {
unsafe fn new(array: &mut [MaybeUninit<T>; N]) -> Self {
Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 }
}
}

impl<T, const N: usize> Drop for FrontGuard<T, N> {
fn drop(&mut self) {
debug_assert!(self.init <= N);
// SAFETY: This raw slice will only contain the initialized objects
// within the buffer.
unsafe {
let slice = ptr::slice_from_raw_parts_mut(self.ptr, self.init);
ptr::drop_in_place(slice);
}
}
}

#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N>
where
Expand All @@ -157,29 +108,20 @@ where
R: Try<Output = B>,
{
// We are iterating from the back we need to first handle the remainder.
if self.next_back_remainder().is_none() {
return R::from_output(init);
}
self.next_back_remainder();

let mut array = MaybeUninit::uninit_array();
// SAFETY: `array` will still be valid if `guard` is dropped.
let mut guard = unsafe { BackGuard::new(&mut array) };
let mut acc = init;
let mut iter = self.iter.by_ref().rev();

self.iter.try_rfold(init, |mut acc, item| {
guard.uninit -= 1;
// SAFETY: `uninit` starts at N, decreases by one each iteration and
// is reset to N once it reaches 0.
unsafe { array.get_unchecked_mut(guard.uninit) }.write(item);
if guard.uninit == 0 {
guard.uninit = N;
let array = mem::replace(&mut array, MaybeUninit::uninit_array());
// SAFETY: the condition above asserts that all elements are
// initialized.
let item = unsafe { MaybeUninit::array_assume_init(array) };
acc = f(acc, item)?;
}
R::from_output(acc)
})
// NB remainder is handled by `next_back_remainder`, so
// `next_chunk` can't return `Err` with non-empty remainder
// (assuming correct `I as ExactSizeIterator` impl).
while let Ok(mut chunk) = iter.next_chunk() {
chunk.reverse();
acc = f(acc, chunk)?
}

try { acc }
}

fn rfold<B, F>(mut self, init: B, mut f: F) -> B
Expand All @@ -195,63 +137,26 @@ impl<I, const N: usize> ArrayChunks<I, N>
where
I: DoubleEndedIterator + ExactSizeIterator,
{
#[inline]
fn next_back_remainder(&mut self) -> Option<()> {
/// Updates `self.remainder` such that `self.iter.len` is divisible by `N`.
fn next_back_remainder(&mut self) {
// Make sure to not override `self.remainder` with an empty array
// when `next_back` is called after `ArrayChunks` exhaustion.
if self.remainder.is_some() {
return;
}

// We use the `ExactSizeIterator` implementation of the underlying
// iterator to know how many remaining elements there are.
let rem = self.iter.len() % N;
if rem == 0 {
return Some(());
}

let mut array = MaybeUninit::uninit_array();

// SAFETY: The array will still be valid if `guard` is dropped and
// it is forgotten otherwise.
let mut guard = unsafe { FrontGuard::new(&mut array) };
// Take the last `rem` elements out of `self.iter`.
let mut remainder =
// SAFETY: `unwrap_err` always succeeds because x % N < N for all x.
unsafe { self.iter.by_ref().rev().take(rem).next_chunk().unwrap_err_unchecked() };

// SAFETY: `rem` is in the range 1..N based on how it is calculated.
for slot in unsafe { array.get_unchecked_mut(..rem) }.iter_mut() {
slot.write(self.iter.next_back()?);
guard.init += 1;
}

let init = guard.init;
mem::forget(guard);
// SAFETY: `array` was initialized with exactly `init` elements.
self.remainder = unsafe {
array.get_unchecked_mut(..init).reverse();
Some(array::IntoIter::new_unchecked(array, 0..init))
};
Some(())
}
}

/// A guard for an array where elements are filled from the right.
struct BackGuard<T, const N: usize> {
/// A pointer to the array that is being filled. We need to use a raw
/// pointer here because of the lifetime issues in the rfold implementations.
ptr: *mut T,
/// The number of *uninitialized* elements.
uninit: usize,
}

impl<T, const N: usize> BackGuard<T, N> {
unsafe fn new(array: &mut [MaybeUninit<T>; N]) -> Self {
Self { ptr: MaybeUninit::slice_as_mut_ptr(array), uninit: N }
}
}

impl<T, const N: usize> Drop for BackGuard<T, N> {
fn drop(&mut self) {
debug_assert!(self.uninit <= N);
// SAFETY: This raw slice will only contain the initialized objects
// within the buffer.
unsafe {
let ptr = self.ptr.offset(self.uninit as isize);
let slice = ptr::slice_from_raw_parts_mut(ptr, N - self.uninit);
ptr::drop_in_place(slice);
}
// We used `.rev()` above, so we need to re-reverse the reminder
remainder.as_mut_slice().reverse();
self.remainder = Some(remainder);
}
}

Expand Down

0 comments on commit 69fc19a

Please sign in to comment.