Skip to content

Commit

Permalink
Add Iterator::array_windows()
Browse files Browse the repository at this point in the history
  • Loading branch information
rossmacarthur committed Dec 29, 2021
1 parent 442248d commit 6f19676
Show file tree
Hide file tree
Showing 7 changed files with 222 additions and 3 deletions.
102 changes: 102 additions & 0 deletions library/core/src/iter/adapters/array_windows.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
use crate::iter::{ExactSizeIterator, Fuse, FusedIterator, Iterator, TrustedLen};

use crate::array;

/// An iterator over all contiguous windows of length `N`. The windows overlap.
/// If the iterator is shorter than `N`, the iterator returns no values.
///
/// This `struct` is created by [`Iterator::array_windows`]. See its
/// documentation for more.
#[derive(Debug, Clone)]
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[unstable(feature = "iter_array_windows", reason = "recently added", issue = "none")]
pub struct ArrayWindows<I, const N: usize>
where
I: Iterator,
I::Item: Clone,
{
iter: Fuse<I>,
last: Option<[I::Item; N]>,
}

impl<I, const N: usize> ArrayWindows<I, N>
where
I: Iterator,
I::Item: Clone,
{
#[inline]
pub(in crate::iter) fn new(iter: I) -> Self {
assert!(N != 0);
Self { iter: iter.fuse(), last: None }
}
}

#[unstable(feature = "iter_array_windows", reason = "recently added", issue = "none")]
impl<I: Iterator, const N: usize> Iterator for ArrayWindows<I, N>
where
I: Iterator,
I::Item: Clone,
{
type Item = [I::Item; N];

#[inline]
fn next(&mut self) -> Option<Self::Item> {
let Self { iter, last } = self;

match last {
Some(last) => {
let item = iter.next()?;
last.rotate_left(1);
if let Some(end) = last.last_mut() {
*end = item;
}
Some(last.clone())
}
None => {
let tmp = array::try_from_fn(|_| iter.next())?;
*last = Some(tmp.clone());
Some(tmp)
}
}
}

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let (lower, upper) = self.iter.size_hint();
// Keep infinite iterator size hint lower bound as `usize::MAX`
if lower == usize::MAX {
(lower, upper)
} else {
(lower.saturating_sub(N - 1), upper.map(|n| n.saturating_sub(N - 1)))
}
}

#[inline]
fn count(self) -> usize {
self.iter.count().saturating_sub(N - 1)
}
}

#[unstable(feature = "iter_array_windows", reason = "recently added", issue = "none")]
impl<I, const N: usize> FusedIterator for ArrayWindows<I, N>
where
I: FusedIterator,
I::Item: Clone,
{
}

#[unstable(feature = "iter_array_windows", reason = "recently added", issue = "none")]
impl<I, const N: usize> ExactSizeIterator for ArrayWindows<I, N>
where
I: ExactSizeIterator,
I::Item: Clone,
{
}

#[unstable(feature = "trusted_len", issue = "37572")]
unsafe impl<I, const N: usize> TrustedLen for ArrayWindows<I, N>
where
I: TrustedLen,
I::Item: Clone,
{
}
4 changes: 4 additions & 0 deletions library/core/src/iter/adapters/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use crate::iter::{InPlaceIterable, Iterator};
use crate::ops::{ControlFlow, Try};

mod array_windows;
mod chain;
mod cloned;
mod copied;
Expand Down Expand Up @@ -30,6 +31,9 @@ pub use self::{
scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip,
};

#[unstable(feature = "iter_array_windows", reason = "recently added", issue = "none")]
pub use self::array_windows::ArrayWindows;

#[stable(feature = "iter_cloned", since = "1.1.0")]
pub use self::cloned::Cloned;

Expand Down
2 changes: 2 additions & 0 deletions library/core/src/iter/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -393,6 +393,8 @@ pub use self::traits::{

#[stable(feature = "iter_zip", since = "1.59.0")]
pub use self::adapters::zip;
#[unstable(feature = "iter_array_windows", reason = "recently added", issue = "none")]
pub use self::adapters::ArrayWindows;
#[stable(feature = "iter_cloned", since = "1.1.0")]
pub use self::adapters::Cloned;
#[stable(feature = "iter_copied", since = "1.36.0")]
Expand Down
50 changes: 47 additions & 3 deletions library/core/src/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,13 @@ use crate::cmp::{self, Ordering};
use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, Residual, Try};

use super::super::TrustedRandomAccessNoCoerce;
use super::super::{
ArrayWindows, Inspect, Map, MapWhile, Peekable, Rev, Scan, Skip, SkipWhile, StepBy, Take,
TakeWhile,
};
use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
use super::super::{FlatMap, Flatten};
use super::super::{FromIterator, Intersperse, IntersperseWith, Product, Sum, Zip};
use super::super::{
Inspect, Map, MapWhile, Peekable, Rev, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile,
};

fn _assert_is_object_safe(_: &dyn Iterator<Item = ()>) {}

Expand Down Expand Up @@ -3057,6 +3058,49 @@ pub trait Iterator {
Cycle::new(self)
}

/// Returns an iterator over all contiguous windows of length `N`. The
/// windows overlap. If the iterator is shorter than `N`, the iterator
/// returns no values.
///
/// `array_windows` clones the iterator elements so that they can be part of
/// successive windows, this makes this it most suited for iterators of
/// references and other values that are cheap to clone.
///
/// # Panics
///
/// If called with `N = 0`.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(iter_array_windows)]
///
/// let mut iter = "rust".chars().array_windows();
/// assert_eq!(iter.next(), Some(['r', 'u']));
/// assert_eq!(iter.next(), Some(['u', 's']));
/// assert_eq!(iter.next(), Some(['s', 't']));
/// assert_eq!(iter.next(), None);
/// ```
///
/// ```
/// #![feature(iter_array_windows)]
///
/// let seq: &[i32] = &[0, 1, 1, 2, 3, 5, 8, 13];
/// for [x, y, z] in seq.iter().copied().array_windows() {
/// assert_eq!(x + y, z);
/// }
/// ```
#[unstable(feature = "iter_array_windows", reason = "recently added", issue = "none")]
fn array_windows<const N: usize>(self) -> ArrayWindows<Self, N>
where
Self: Sized,
Self::Item: Clone,
{
ArrayWindows::new(self)
}

/// Sums the elements of an iterator.
///
/// Takes each element, adds them together, and returns the result.
Expand Down
65 changes: 65 additions & 0 deletions library/core/tests/iter/adapters/array_windows.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
use core::iter;

#[test]
fn test_array_windows_infer() {
let s = [0, 1, 0, 1, 0, 1];
for [a, b] in s.iter().copied().array_windows() {
assert_eq!(a + b, 1);
}
for [a, b, c, d] in s.iter().copied().array_windows() {
assert_eq!(a + b + c + d, 2);
}
}

#[test]
fn test_array_windows_size_hint() {
let iter = (0..6).array_windows::<1>();
assert_eq!(iter.size_hint(), (6, Some(6)));

let iter = (0..6).array_windows::<3>();
assert_eq!(iter.size_hint(), (4, Some(4)));

let iter = (0..6).array_windows::<5>();
assert_eq!(iter.size_hint(), (2, Some(2)));

let iter = (0..6).array_windows::<7>();
assert_eq!(iter.size_hint(), (0, Some(0)));

let iter = (1..).array_windows::<2>();
assert_eq!(iter.size_hint(), (usize::MAX, None));

let iter = (1..).filter(|x| x % 2 != 0).array_windows::<2>();
assert_eq!(iter.size_hint(), (0, None));
}

#[test]
fn test_array_windows_count() {
let iter = (0..6).array_windows::<1>();
assert_eq!(iter.count(), 6);

let iter = (0..6).array_windows::<3>();
assert_eq!(iter.count(), 4);

let iter = (0..6).array_windows::<5>();
assert_eq!(iter.count(), 2);

let iter = (0..6).array_windows::<7>();
assert_eq!(iter.count(), 0);

let iter = (0..6).filter(|x| x % 2 == 0).array_windows::<2>();
assert_eq!(iter.count(), 2);

let iter = iter::empty::<i32>().array_windows::<2>();
assert_eq!(iter.count(), 0);

let iter = [(); usize::MAX].iter().array_windows::<2>();
assert_eq!(iter.count(), usize::MAX - 1);
}

#[test]
fn test_array_windows_nth() {
let mut iter = (0..6).array_windows::<4>();
assert_eq!(iter.nth(1), Some([1, 2, 3, 4]));
assert_eq!(iter.nth(0), Some([2, 3, 4, 5]));
assert_eq!(iter.nth(1), None);
}
1 change: 1 addition & 0 deletions library/core/tests/iter/adapters/mod.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
mod array_windows;
mod chain;
mod cloned;
mod copied;
Expand Down
1 change: 1 addition & 0 deletions library/core/tests/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@
#![feature(slice_partition_dedup)]
#![feature(int_log)]
#![feature(iter_advance_by)]
#![feature(iter_array_windows)]
#![feature(iter_partition_in_place)]
#![feature(iter_intersperse)]
#![feature(iter_is_partitioned)]
Expand Down

0 comments on commit 6f19676

Please sign in to comment.