From 8e6fcfd7f4c2d02e0ca572e036b09df69b1a88a7 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Tue, 25 May 2021 20:57:41 +0100 Subject: [PATCH 1/2] feat: arrays and next_array creates an ArrayVec type that handles uninitialised array data --- src/array_impl.rs | 144 ++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 57 ++++++++++++++++++ tests/arrays.rs | 39 +++++++++++++ tests/quick.rs | 26 +++++++++ 4 files changed, 266 insertions(+) create mode 100644 src/array_impl.rs create mode 100644 tests/arrays.rs diff --git a/src/array_impl.rs b/src/array_impl.rs new file mode 100644 index 000000000..f12fd8a3a --- /dev/null +++ b/src/array_impl.rs @@ -0,0 +1,144 @@ +/// An iterator that groups the items in arrays of a specific size. +/// +/// See [`.arrays()`](crate::Itertools::arrays) for more information. +#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] +pub struct Arrays +{ + iter: I, + buf: ArrayVec, +} + +impl Arrays { + pub fn new(iter: I) -> Self { + Self { iter, buf: ArrayVec::new() } + } +} +impl Arrays where I::Item: Clone { + pub fn remaining(self) -> Vec { + self.buf.into_vec() + } +} + +impl Iterator for Arrays { + type Item = [I::Item; N]; + + fn next(&mut self) -> Option { + // SAFETY: + // ArrayVec::push_unchecked is safe as long as len < N. + // This is guaranteed by the for loop + unsafe { + for _ in self.buf.len()..N { + self.buf.push_unchecked(self.iter.next()?) + } + + Some(self.buf.take_unchecked()) + } + } +} + +pub fn next_array(iter: &mut I) -> Option<[I::Item; N]> { + let mut array_vec = ArrayVec::new(); + + // SAFETY: + // ArrayVec::push_unchecked is safe as long as len < N. + // This is guaranteed by the for loop + unsafe { + for _ in 0..N { + array_vec.push_unchecked(iter.next()?) + } + } + + array_vec.into_array() +} + +// ArrayVec is a safe wrapper around a [T; N]. +// It allows safely initialising an empty array +pub struct ArrayVec { + data: std::mem::MaybeUninit<[T; N]>, + len: usize, +} + +impl Drop for ArrayVec { + fn drop(&mut self) { + // SAFETY: + // The contract of the ArrayVec ensures that data[..len] is initialised + unsafe { + let ptr = self.data.as_mut_ptr() as *mut T; + drop(std::slice::from_raw_parts_mut(ptr, self.len)); + } + } +} + +impl ArrayVec { + pub const fn new() -> Self { + Self { + data: std::mem::MaybeUninit::uninit(), + len: 0, + } + } + + pub fn push(&mut self, v: T) { + assert!(self.len < N); + // SAFETY: asserts that len < N. + unsafe { self.push_unchecked(v) } + } + + // Unsafety: + // len must be less than N. If `len < N` is guaranteed, this operation is safe + // This is because the contract of ArrayVec guarantees that if len < N, then the value + // at len is valid and uninitialised. + pub unsafe fn push_unchecked(&mut self, v: T) { + // The contract of ArrayVec guarantees that the value at self.len, if < N, + // is uninitialised, and therefore does not need dropping. So use write to + // overwrite the value + let ptr = (self.data.as_mut_ptr() as *mut T).add(self.len); + std::ptr::write(ptr, v); + self.len += 1; + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn into_array(self) -> Option<[T; N]> { + if self.len == N { + // SAFETY: + // If len == N, then all the data is initialised and this is safe + unsafe { Some(self.into_array_unchecked()) } + } else { + None + } + } + + // Unsafety: + // len must be equal to N. If `len == N` is guaranteed, this operation is safe. + // This is because the contract of ArrayVec guarantees that if len == N, all the values + // have been initialised correctly. + unsafe fn into_array_unchecked(mut self) -> [T; N] { + // move out without dropping + let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init()); + std::mem::forget(self); + data.assume_init() + } + + // Unsafety: + // len must be equal to N. If `len == N` is guaranteed, this operation is safe. + // This is because the contract of ArrayVec guarantees that if len == N, all the values + // have been initialised correctly. + unsafe fn take_unchecked(&mut self) -> [T; N] { + let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init()); + self.len = 0; + data.assume_init() + } +} + +impl ArrayVec { + pub fn into_vec(mut self) -> Vec { + unsafe { + let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init()); + let len = self.len; + std::mem::forget(self); + std::slice::from_raw_parts(data.as_ptr() as *const T, len).to_vec() + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 1a6632e20..1db6d8ca9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -112,6 +112,7 @@ pub mod structs { pub use crate::adaptors::{MapResults, Step}; #[cfg(feature = "use_alloc")] pub use crate::adaptors::MultiProduct; + pub use crate::array_impl::Arrays; #[cfg(feature = "use_alloc")] pub use crate::combinations::Combinations; #[cfg(feature = "use_alloc")] @@ -181,6 +182,7 @@ pub use crate::sources::{repeat_call, unfold, iterate}; pub use crate::with_position::Position; pub use crate::ziptuple::multizip; mod adaptors; +mod array_impl; mod either_or_both; pub use crate::either_or_both::EitherOrBoth; #[doc(hidden)] @@ -775,6 +777,42 @@ pub trait Itertools : Iterator { tuple_impl::tuples(self) } + /// Return an iterator that groups the items in arrays of a specific size + /// + /// See also the method [`.next_array()`](#method.next_array). + /// + /// ``` + /// use itertools::Itertools; + /// let mut v = Vec::new(); + /// for [a, b] in (1..5).arrays() { + /// v.push([a, b]); + /// } + /// assert_eq!(v, vec![[1, 2], [3, 4]]); + /// + /// let mut it = (1..8).arrays(); + /// assert_eq!(Some([1, 2, 3]), it.next()); + /// assert_eq!(Some([4, 5, 6]), it.next()); + /// assert_eq!(None, it.next()); + /// // get any remaining data + /// assert_eq!(vec![7], it.remaining()); + /// + /// // this requires a type hint + /// let it = (1..7).arrays::<3>(); + /// itertools::assert_equal(it, vec![[1, 2, 3], [4, 5, 6]]); + /// + /// // you can also specify the complete type + /// use itertools::Arrays; + /// use std::ops::Range; + /// + /// let it: Arrays, 3> = (1..7).arrays(); + /// itertools::assert_equal(it, vec![[1, 2, 3], [4, 5, 6]]); + /// ``` + fn arrays(self) -> array_impl::Arrays + where Self: Sized + Iterator, + { + array_impl::Arrays::new(self) + } + /// Split into an iterator pair that both yield all elements from /// the original iterator. /// @@ -1702,6 +1740,25 @@ pub trait Itertools : Iterator { T::collect_from_iter_no_buf(self) } + /// Advances the iterator and returns the next items grouped in an array of + /// a specific size. + /// + /// If there are enough elements to be grouped in a tuple, then the tuple is + /// returned inside `Some`, otherwise `None` is returned. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let mut iter = 1..5; + /// + /// assert_eq!(Some([1, 2]), iter.next_array()); + /// ``` + fn next_array(&mut self) -> Option<[Self::Item; N]> + where Self: Sized + Iterator + { + array_impl::next_array(self) + } + /// Collects all items from the iterator into a tuple of a specific size /// (up to 12). /// diff --git a/tests/arrays.rs b/tests/arrays.rs new file mode 100644 index 000000000..f60037e88 --- /dev/null +++ b/tests/arrays.rs @@ -0,0 +1,39 @@ +use itertools::Itertools; + +#[test] +fn arrays() { + let v = [1, 2, 3, 4, 5]; + let mut iter = v.iter().cloned().arrays(); + assert_eq!(Some([1]), iter.next()); + assert_eq!(Some([2]), iter.next()); + assert_eq!(Some([3]), iter.next()); + assert_eq!(Some([4]), iter.next()); + assert_eq!(Some([5]), iter.next()); + assert_eq!(None, iter.next()); + assert_eq!(Vec::::new(), iter.remaining()); + + let mut iter = v.iter().cloned().arrays(); + assert_eq!(Some([1, 2]), iter.next()); + assert_eq!(Some([3, 4]), iter.next()); + assert_eq!(None, iter.next()); + assert_eq!(vec![5], iter.remaining()); + + let mut iter = v.iter().cloned().arrays(); + assert_eq!(Some([1, 2, 3]), iter.next()); + assert_eq!(None, iter.next()); + assert_eq!(vec![4, 5], iter.remaining()); + + let mut iter = v.iter().cloned().arrays(); + assert_eq!(Some([1, 2, 3, 4]), iter.next()); + assert_eq!(None, iter.next()); + assert_eq!(vec![5], iter.remaining()); +} + +#[test] +fn next_array() { + let v = [1, 2, 3, 4, 5]; + let mut iter = v.iter(); + assert_eq!(iter.next_array().map(|[&x, &y]| (x, y)), Some((1, 2))); + assert_eq!(iter.next_array().map(|[&x, &y]| (x, y)), Some((3, 4))); + assert_eq!(iter.next_array::<2>(), None); +} diff --git a/tests/quick.rs b/tests/quick.rs index 7769cb432..be84a1327 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -1073,6 +1073,32 @@ quickcheck! { } } +quickcheck! { + fn equal_arrays_1(a: Vec) -> bool { + let x = a.chunks(1).map(|s| [&s[0]]); + let y = a.iter().arrays::<1>(); + itertools::equal(x, y) + } + + fn equal_arrays_2(a: Vec) -> bool { + let x = a.chunks(2).filter(|s| s.len() == 2).map(|s| [&s[0], &s[1]]); + let y = a.iter().arrays::<2>(); + itertools::equal(x, y) + } + + fn equal_arrays_3(a: Vec) -> bool { + let x = a.chunks(3).filter(|s| s.len() == 3).map(|s| [&s[0], &s[1], &s[2]]); + let y = a.iter().arrays::<3>(); + itertools::equal(x, y) + } + + fn equal_arrays_4(a: Vec) -> bool { + let x = a.chunks(4).filter(|s| s.len() == 4).map(|s| [&s[0], &s[1], &s[2], &s[3]]); + let y = a.iter().arrays::<4>(); + itertools::equal(x, y) + } +} + // with_position quickcheck! { fn with_position_exact_size_1(a: Vec) -> bool { From 8270ed2530922ad68037b4513f4cd0feba305a05 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Tue, 25 May 2021 21:18:25 +0100 Subject: [PATCH 2/2] feat: more docs, more comments --- src/array_impl.rs | 94 ++++++++++++++++++++++++++++++----------------- 1 file changed, 61 insertions(+), 33 deletions(-) diff --git a/src/array_impl.rs b/src/array_impl.rs index f12fd8a3a..3243a497c 100644 --- a/src/array_impl.rs +++ b/src/array_impl.rs @@ -9,13 +9,18 @@ pub struct Arrays } impl Arrays { + /// Creates a new Arrays iterator. + /// See [`.arrays()`](crate::Itertools::arrays) for more information. pub fn new(iter: I) -> Self { Self { iter, buf: ArrayVec::new() } } } impl Arrays where I::Item: Clone { - pub fn remaining(self) -> Vec { - self.buf.into_vec() + /// Returns any remaining data left in the iterator. + /// This is useful if the iterator has left over values + /// that didn't fit into the array size. + pub fn remaining(&self) -> &[I::Item] { + self.buf.into_slice() } } @@ -36,19 +41,23 @@ impl Iterator for Arrays { } } +/// Reads N elements from the iter, returning them in an array. If there's not enough +/// elements, returns None. pub fn next_array(iter: &mut I) -> Option<[I::Item; N]> { let mut array_vec = ArrayVec::new(); - // SAFETY: - // ArrayVec::push_unchecked is safe as long as len < N. - // This is guaranteed by the for loop unsafe { + // SAFETY: + // ArrayVec::push_unchecked is safe as long as len < N. + // This is guaranteed by the for loop for _ in 0..N { array_vec.push_unchecked(iter.next()?) } - } - array_vec.into_array() + // SAFETY: + // We have guaranteed to have filled all N elements + Some(array_vec.into_array_unchecked()) + } } // ArrayVec is a safe wrapper around a [T; N]. @@ -70,6 +79,7 @@ impl Drop for ArrayVec { } impl ArrayVec { + /// Creates a new empty ArrayVec pub const fn new() -> Self { Self { data: std::mem::MaybeUninit::uninit(), @@ -77,16 +87,36 @@ impl ArrayVec { } } + /// Returns the number of initialised elements in the ArrayVec + pub fn len(&self) -> usize { + self.len + } + + /// Returns whether the ArrayVec is empty + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Returns whether the ArrayVec is full + pub fn is_full(&self) -> bool { + self.len == N + } + + /// Push a value into the ArrayVec + /// + /// Panics: + /// If the ArrayVec is full, this function will panic pub fn push(&mut self, v: T) { - assert!(self.len < N); - // SAFETY: asserts that len < N. + assert!(!self.is_full()); + // SAFETY: asserted that self is not full unsafe { self.push_unchecked(v) } } - // Unsafety: - // len must be less than N. If `len < N` is guaranteed, this operation is safe - // This is because the contract of ArrayVec guarantees that if len < N, then the value - // at len is valid and uninitialised. + /// Push a value into the ArrayVec + /// + /// Unsafety: + /// The ArrayVec must not be full. If the ArrayVec is full, this function will try write data + /// out-of-bounds. pub unsafe fn push_unchecked(&mut self, v: T) { // The contract of ArrayVec guarantees that the value at self.len, if < N, // is uninitialised, and therefore does not need dropping. So use write to @@ -96,12 +126,9 @@ impl ArrayVec { self.len += 1; } - pub fn len(&self) -> usize { - self.len - } - + /// If the ArrayVec is full, returns the data owned. Otherwise, returns None pub fn into_array(self) -> Option<[T; N]> { - if self.len == N { + if self.is_full() { // SAFETY: // If len == N, then all the data is initialised and this is safe unsafe { Some(self.into_array_unchecked()) } @@ -110,10 +137,11 @@ impl ArrayVec { } } - // Unsafety: - // len must be equal to N. If `len == N` is guaranteed, this operation is safe. - // This is because the contract of ArrayVec guarantees that if len == N, all the values - // have been initialised correctly. + /// Returns the data owned by the ArrayVec + /// + /// Unsafety: + /// The ArrayVec must be full. If it is not full, some of the values in the array will be + /// unintialised. This will cause undefined behaviour unsafe fn into_array_unchecked(mut self) -> [T; N] { // move out without dropping let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init()); @@ -121,24 +149,24 @@ impl ArrayVec { data.assume_init() } - // Unsafety: - // len must be equal to N. If `len == N` is guaranteed, this operation is safe. - // This is because the contract of ArrayVec guarantees that if len == N, all the values - // have been initialised correctly. + /// Returns the data owned by the ArrayVec, resetting the ArrayVec to an empty state + /// + /// Unsafety: + /// The ArrayVec must be full. If it is not full, some of the values in the array will be + /// unintialised. This will cause undefined behaviour unsafe fn take_unchecked(&mut self) -> [T; N] { let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init()); self.len = 0; data.assume_init() } -} -impl ArrayVec { - pub fn into_vec(mut self) -> Vec { + /// Borrows the initialised data in the ArrayVec + pub fn into_slice(&self) -> &[T] { unsafe { - let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init()); - let len = self.len; - std::mem::forget(self); - std::slice::from_raw_parts(data.as_ptr() as *const T, len).to_vec() + // let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init()); + // let len = self.len; + // std::mem::forget(self); + std::slice::from_raw_parts(self.data.as_ptr() as *const T, self.len) } } }