Use array::IntoIter for the ArrayChunks remainder

This commit is contained in:
Ross MacArthur 2022-02-04 15:57:58 +02:00 committed by Maybe Waffle
parent ca3d1010bb
commit f5485181ca
4 changed files with 33 additions and 105 deletions

View File

@ -84,6 +84,16 @@ impl<T, const N: usize> IntoIter<T, N> {
IntoIterator::into_iter(array) 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. /// Creates an iterator over the elements in a partially-initialized buffer.
/// ///
/// If you have a fully-initialized array, then use [`IntoIterator`]. /// If you have a fully-initialized array, then use [`IntoIterator`].

View File

@ -1,68 +1,10 @@
use crate::array;
use crate::iter::{Fuse, FusedIterator, Iterator, TrustedLen}; use crate::iter::{Fuse, FusedIterator, Iterator, TrustedLen};
use crate::mem; use crate::mem;
use crate::mem::MaybeUninit; use crate::mem::MaybeUninit;
use crate::ops::{ControlFlow, Try}; use crate::ops::{ControlFlow, Try};
use crate::ptr; 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. /// 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 /// The chunks do not overlap. If `N` does not divide the length of the
@ -75,7 +17,7 @@ impl<T, const N: usize> Drop for Remainder<T, N> {
#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
pub struct ArrayChunks<I: Iterator, const N: usize> { pub struct ArrayChunks<I: Iterator, const N: usize> {
iter: Fuse<I>, iter: Fuse<I>,
remainder: Remainder<I::Item, N>, remainder: Option<array::IntoIter<I::Item, N>>,
} }
impl<I, const N: usize> ArrayChunks<I, N> impl<I, const N: usize> ArrayChunks<I, N>
@ -84,25 +26,16 @@ where
{ {
pub(in crate::iter) fn new(iter: I) -> Self { pub(in crate::iter) fn new(iter: I) -> Self {
assert!(N != 0, "chunk size must be non-zero"); assert!(N != 0, "chunk size must be non-zero");
Self { iter: iter.fuse(), remainder: Remainder::new() } Self { iter: iter.fuse(), remainder: None }
} }
/// Returns a reference to the remaining elements of the original iterator /// Returns an iterator over the remaining elements of the original iterator
/// that are not going to be returned by this iterator. The returned slice /// that are not going to be returned by this iterator. The returned
/// has at most `N-1` elements. /// iterator will yield at most `N-1` elements.
#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")] #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
#[inline] #[inline]
pub fn remainder(&self) -> &[I::Item] { pub fn into_remainder(self) -> Option<array::IntoIter<I::Item, N>> {
self.remainder.as_slice() self.remainder
}
/// 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.
#[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()
} }
} }
@ -129,8 +62,10 @@ where
if guard.init > 0 { if guard.init > 0 {
let init = guard.init; let init = guard.init;
mem::forget(guard); mem::forget(guard);
// SAFETY: `array` was initialized with `init` elements. self.remainder = {
self.remainder = unsafe { Remainder::with_init(array, init) }; // SAFETY: `array` was initialized with `init` elements.
Some(unsafe { array::IntoIter::with_partial(array, 0..init) })
};
} }
return None; return None;
} }
@ -189,7 +124,7 @@ where
let init = guard.init; let init = guard.init;
mem::forget(guard); mem::forget(guard);
// SAFETY: `array` was initialized with `init` elements. // 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) R::from_output(o)
} }
@ -370,7 +305,7 @@ where
// SAFETY: `array` was initialized with exactly `init` elements. // SAFETY: `array` was initialized with exactly `init` elements.
self.remainder = unsafe { self.remainder = unsafe {
array.get_unchecked_mut(..init).reverse(); array.get_unchecked_mut(..init).reverse();
Remainder::with_init(array, init) Some(array::IntoIter::with_partial(array, 0..init))
}; };
Some(()) Some(())
} }

View File

@ -3319,7 +3319,9 @@ pub trait Iterator {
/// Returns an iterator over `N` elements of the iterator at a time. /// 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 /// 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 /// # Panics
/// ///
@ -3336,7 +3338,7 @@ pub trait Iterator {
/// assert_eq!(iter.next(), Some(['l', 'o'])); /// assert_eq!(iter.next(), Some(['l', 'o']));
/// assert_eq!(iter.next(), Some(['r', 'e'])); /// assert_eq!(iter.next(), Some(['r', 'e']));
/// assert_eq!(iter.next(), None); /// assert_eq!(iter.next(), None);
/// assert_eq!(iter.remainder(), &['m']); /// assert_eq!(iter.into_remainder().unwrap().as_slice(), &['m']);
/// ``` /// ```
/// ///
/// ``` /// ```

View File

@ -15,39 +15,24 @@ fn test_iterator_array_chunks_infer() {
fn test_iterator_array_chunks_clone_and_drop() { fn test_iterator_array_chunks_clone_and_drop() {
let count = Cell::new(0); let count = Cell::new(0);
let mut it = (0..5).map(|_| CountDrop::new(&count)).array_chunks::<3>(); let mut it = (0..5).map(|_| CountDrop::new(&count)).array_chunks::<3>();
assert_eq!(it.by_ref().count(), 1); assert_eq!(it.by_ref().count(), 1);
assert_eq!(count.get(), 3); assert_eq!(count.get(), 3);
assert_eq!(it.remainder().len(), 2);
let mut it2 = it.clone(); let mut it2 = it.clone();
assert_eq!(count.get(), 3); assert_eq!(count.get(), 3);
assert_eq!(it2.remainder().len(), 2); assert_eq!(it.into_remainder().unwrap().len(), 2);
drop(it);
assert_eq!(count.get(), 5); assert_eq!(count.get(), 5);
assert_eq!(it2.remainder().len(), 2);
assert!(it2.next().is_none()); assert!(it2.next().is_none());
assert_eq!(it2.into_remainder().unwrap().len(), 2);
drop(it2);
assert_eq!(count.get(), 7); assert_eq!(count.get(), 7);
} }
#[test] #[test]
fn test_iterator_array_chunks_remainder() { fn test_iterator_array_chunks_remainder() {
let mut it = (0..11).array_chunks::<4>(); 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.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.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.next(), None); assert_eq!(it.into_remainder().unwrap().as_slice(), &[8, 9, 10]);
assert_eq!(it.remainder(), &[8, 9, 10]);
assert_eq!(it.remainder_mut(), &[8, 9, 10]);
} }
#[test] #[test]
@ -105,8 +90,7 @@ fn test_iterator_array_chunks_next_and_next_back() {
assert_eq!(it.next(), None); assert_eq!(it.next(), None);
assert_eq!(it.next_back(), None); assert_eq!(it.next_back(), None);
assert_eq!(it.next(), None); assert_eq!(it.next(), None);
assert_eq!(it.remainder(), &[9, 10]); assert_eq!(it.into_remainder().unwrap().as_slice(), &[9, 10]);
assert_eq!(it.remainder_mut(), &[9, 10]);
} }
#[test] #[test]
@ -119,7 +103,7 @@ fn test_iterator_array_chunks_rev_remainder() {
assert_eq!(it.next(), None); assert_eq!(it.next(), None);
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] #[test]
@ -128,7 +112,6 @@ fn test_iterator_array_chunks_try_fold() {
let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>(); 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)); let result: Result<_, ()> = it.by_ref().try_fold(0, |acc, _item| Ok(acc + 1));
assert_eq!(result, Ok(3)); assert_eq!(result, Ok(3));
assert_eq!(it.remainder().len(), 1);
assert_eq!(count.get(), 9); assert_eq!(count.get(), 9);
drop(it); drop(it);
assert_eq!(count.get(), 10); assert_eq!(count.get(), 10);
@ -137,7 +120,6 @@ fn test_iterator_array_chunks_try_fold() {
let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>(); 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) }); 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!(result, Err(2));
assert_eq!(it.remainder().len(), 0);
assert_eq!(count.get(), 9); assert_eq!(count.get(), 9);
drop(it); drop(it);
assert_eq!(count.get(), 9); assert_eq!(count.get(), 9);
@ -166,7 +148,6 @@ fn test_iterator_array_chunks_try_rfold() {
let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>(); let mut it = (0..10).map(|_| CountDrop::new(&count)).array_chunks::<3>();
let result: Result<_, ()> = it.try_rfold(0, |acc, _item| Ok(acc + 1)); let result: Result<_, ()> = it.try_rfold(0, |acc, _item| Ok(acc + 1));
assert_eq!(result, Ok(3)); assert_eq!(result, Ok(3));
assert_eq!(it.remainder().len(), 1);
assert_eq!(count.get(), 9); assert_eq!(count.get(), 9);
drop(it); drop(it);
assert_eq!(count.get(), 10); assert_eq!(count.get(), 10);