Skip to content

Commit

Permalink
Use array::IntoIter for the ArrayChunks remainder
Browse files Browse the repository at this point in the history
  • Loading branch information
rossmacarthur committed Feb 4, 2022
1 parent c115d5a commit 6adcad8
Show file tree
Hide file tree
Showing 4 changed files with 33 additions and 105 deletions.
10 changes: 10 additions & 0 deletions library/core/src/array/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,16 @@ impl<T, const N: usize> IntoIter<T, N> {
IntoIterator::into_iter(array)
}

/// Creates a new iterator from a partially initalized array.
///
/// # Safety
///
/// The caller must guarantee that all and only the `alive` elements of
/// `data` are initialized.
pub(crate) unsafe fn with_partial(data: [MaybeUninit<T>; N], alive: Range<usize>) -> Self {
Self { data, alive }
}

/// Creates an iterator over the elements in a partially-initialized buffer.
///
/// If you have a fully-initialized array, then use [`IntoIterator`].
Expand Down
93 changes: 14 additions & 79 deletions library/core/src/iter/adapters/array_chunks.rs
Original file line number Diff line number Diff line change
@@ -1,68 +1,10 @@
use crate::array;
use crate::iter::{Fuse, FusedIterator, Iterator, TrustedLen};
use crate::mem;
use crate::mem::MaybeUninit;
use crate::ops::{ControlFlow, Try};
use crate::ptr;

#[derive(Debug)]
struct Remainder<T, const N: usize> {
array: [MaybeUninit<T>; N],
init: usize,
}

impl<T, const N: usize> Remainder<T, N> {
fn new() -> Self {
Self { array: MaybeUninit::uninit_array(), init: 0 }
}

unsafe fn with_init(array: [MaybeUninit<T>; N], init: usize) -> Self {
Self { array, init }
}

fn as_slice(&self) -> &[T] {
debug_assert!(self.init <= N);
// SAFETY: This raw slice will only contain the initialized objects
// within the buffer.
unsafe {
let slice = self.array.get_unchecked(..self.init);
MaybeUninit::slice_assume_init_ref(slice)
}
}

fn as_mut_slice(&mut self) -> &mut [T] {
debug_assert!(self.init <= N);
// SAFETY: This raw slice will only contain the initialized objects
// within the buffer.
unsafe {
let slice = self.array.get_unchecked_mut(..self.init);
MaybeUninit::slice_assume_init_mut(slice)
}
}
}

impl<T, const N: usize> Clone for Remainder<T, N>
where
T: Clone,
{
fn clone(&self) -> Self {
let mut new = Self::new();
// SAFETY: The new array is the same size and `init` is always less than
// or equal to `N`.
let this = unsafe { new.array.get_unchecked_mut(..self.init) };
MaybeUninit::write_slice_cloned(this, self.as_slice());
new.init = self.init;
new
}
}

impl<T, const N: usize> Drop for Remainder<T, N> {
fn drop(&mut self) {
// SAFETY: This raw slice will only contain the initialized objects
// within the buffer.
unsafe { ptr::drop_in_place(self.as_mut_slice()) }
}
}

/// An iterator over `N` elements of the iterator at a time.
///
/// The chunks do not overlap. If `N` does not divide the length of the
Expand All @@ -75,7 +17,7 @@ impl<T, const N: usize> Drop for Remainder<T, N> {
#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
pub struct ArrayChunks<I: Iterator, const N: usize> {
iter: Fuse<I>,
remainder: Remainder<I::Item, N>,
remainder: Option<array::IntoIter<I::Item, N>>,
}

impl<I, const N: usize> ArrayChunks<I, N>
Expand All @@ -84,25 +26,16 @@ where
{
pub(in crate::iter) fn new(iter: I) -> Self {
assert!(N != 0, "chunk size must be non-zero");
Self { iter: iter.fuse(), remainder: Remainder::new() }
}

/// Returns a reference to the remaining elements of the original iterator
/// that are not going to be returned by this iterator. The returned slice
/// has at most `N-1` elements.
#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
#[inline]
pub fn remainder(&self) -> &[I::Item] {
self.remainder.as_slice()
Self { iter: iter.fuse(), remainder: None }
}

/// Returns a mutable reference to the remaining elements of the original
/// iterator that are not going to be returned by this iterator. The
/// returned slice has at most `N-1` elements.
/// Returns an iterator over the remaining elements of the original iterator
/// that are not going to be returned by this iterator. The returned
/// iterator will yield at most `N-1` elements.
#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
#[inline]
pub fn remainder_mut(&mut self) -> &mut [I::Item] {
self.remainder.as_mut_slice()
pub fn into_remainder(self) -> Option<array::IntoIter<I::Item, N>> {
self.remainder
}
}

Expand All @@ -129,8 +62,10 @@ where
if guard.init > 0 {
let init = guard.init;
mem::forget(guard);
// SAFETY: `array` was initialized with `init` elements.
self.remainder = unsafe { Remainder::with_init(array, init) };
self.remainder = {
// SAFETY: `array` was initialized with `init` elements.
Some(unsafe { array::IntoIter::with_partial(array, 0..init) })
};
}
return None;
}
Expand Down Expand Up @@ -189,7 +124,7 @@ where
let init = guard.init;
mem::forget(guard);
// SAFETY: `array` was initialized with `init` elements.
self.remainder = unsafe { Remainder::with_init(array, init) };
self.remainder = Some(unsafe { array::IntoIter::with_partial(array, 0..init) });
}
R::from_output(o)
}
Expand Down Expand Up @@ -370,7 +305,7 @@ where
// SAFETY: `array` was initialized with exactly `init` elements.
self.remainder = unsafe {
array.get_unchecked_mut(..init).reverse();
Remainder::with_init(array, init)
Some(array::IntoIter::with_partial(array, 0..init))
};
Some(())
}
Expand Down
6 changes: 4 additions & 2 deletions library/core/src/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3092,7 +3092,9 @@ pub trait Iterator {
/// Returns an iterator over `N` elements of the iterator at a time.
///
/// The chunks do not overlap. If `N` does not divide the length of the
/// iterator, then the last up to `N-1` elements will be omitted.
/// iterator, then the last up to `N-1` elements will be omitted and can be
/// retrieved from the [`.into_remainder()`][ArrayChunks::into_remainder]
/// function of the iterator.
///
/// # Panics
///
Expand All @@ -3109,7 +3111,7 @@ pub trait Iterator {
/// assert_eq!(iter.next(), Some(['l', 'o']));
/// assert_eq!(iter.next(), Some(['r', 'e']));
/// assert_eq!(iter.next(), None);
/// assert_eq!(iter.remainder(), &['m']);
/// assert_eq!(iter.into_remainder().unwrap().as_slice(), &['m']);
/// ```
///
/// ```
Expand Down
29 changes: 5 additions & 24 deletions library/core/tests/iter/adapters/array_chunks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,39 +15,24 @@ fn test_iterator_array_chunks_infer() {
fn test_iterator_array_chunks_clone_and_drop() {
let count = Cell::new(0);
let mut it = (0..5).map(|_| CountDrop::new(&count)).array_chunks::<3>();

assert_eq!(it.by_ref().count(), 1);
assert_eq!(count.get(), 3);
assert_eq!(it.remainder().len(), 2);

let mut it2 = it.clone();
assert_eq!(count.get(), 3);
assert_eq!(it2.remainder().len(), 2);

drop(it);
assert_eq!(it.into_remainder().unwrap().len(), 2);
assert_eq!(count.get(), 5);
assert_eq!(it2.remainder().len(), 2);
assert!(it2.next().is_none());

drop(it2);
assert_eq!(it2.into_remainder().unwrap().len(), 2);
assert_eq!(count.get(), 7);
}

#[test]
fn test_iterator_array_chunks_remainder() {
let mut it = (0..11).array_chunks::<4>();
assert_eq!(it.remainder(), &[]);
assert_eq!(it.remainder_mut(), &[]);
assert_eq!(it.next(), Some([0, 1, 2, 3]));
assert_eq!(it.remainder(), &[]);
assert_eq!(it.remainder_mut(), &[]);
assert_eq!(it.next(), Some([4, 5, 6, 7]));
assert_eq!(it.remainder(), &[]);
assert_eq!(it.remainder_mut(), &[]);
assert_eq!(it.next(), None);
assert_eq!(it.next(), None);
assert_eq!(it.remainder(), &[8, 9, 10]);
assert_eq!(it.remainder_mut(), &[8, 9, 10]);
assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]);
}

#[test]
Expand Down Expand Up @@ -105,8 +90,7 @@ fn test_iterator_array_chunks_next_and_next_back() {
assert_eq!(it.next(), None);
assert_eq!(it.next_back(), None);
assert_eq!(it.next(), None);
assert_eq!(it.remainder(), &[9, 10]);
assert_eq!(it.remainder_mut(), &[9, 10]);
assert_eq!(it.into_remainder().unwrap().as_slice(), &[9, 10]);
}

#[test]
Expand All @@ -119,7 +103,7 @@ fn test_iterator_array_chunks_rev_remainder() {
assert_eq!(it.next(), None);
assert_eq!(it.next(), None);
}
assert_eq!(it.remainder(), &[8, 9, 10]);
assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]);
}

#[test]
Expand All @@ -128,7 +112,6 @@ fn test_iterator_array_chunks_try_fold() {
let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
let result: Result<_, ()> = it.by_ref().try_fold(0, |acc, _item| Ok(acc + 1));
assert_eq!(result, Ok(3));
assert_eq!(it.remainder().len(), 1);
assert_eq!(count.get(), 9);
drop(it);
assert_eq!(count.get(), 10);
Expand All @@ -137,7 +120,6 @@ fn test_iterator_array_chunks_try_fold() {
let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
let result = it.by_ref().try_fold(0, |acc, _item| if acc < 2 { Ok(acc + 1) } else { Err(acc) });
assert_eq!(result, Err(2));
assert_eq!(it.remainder().len(), 0);
assert_eq!(count.get(), 9);
drop(it);
assert_eq!(count.get(), 9);
Expand Down Expand Up @@ -166,7 +148,6 @@ fn test_iterator_array_chunks_try_rfold() {
let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
let result: Result<_, ()> = it.try_rfold(0, |acc, _item| Ok(acc + 1));
assert_eq!(result, Ok(3));
assert_eq!(it.remainder().len(), 1);
assert_eq!(count.get(), 9);
drop(it);
assert_eq!(count.get(), 10);
Expand Down

0 comments on commit 6adcad8

Please sign in to comment.