Implement DoubleEndedIterator for iter::{StepBy, Peekable, Take}

This commit is contained in:
Tim Vermeulen 2019-06-02 13:31:49 +02:00
parent 538e17a3fd
commit 56ebfb185b
2 changed files with 296 additions and 13 deletions

View File

@ -485,6 +485,39 @@ impl<I> Iterator for StepBy<I> where I: Iterator {
}
}
impl<I> StepBy<I> where I: ExactSizeIterator {
// The zero-based index starting from the end of the iterator of the
// last element. Used in the `DoubleEndedIterator` implementation.
fn next_back_index(&self) -> usize {
let rem = self.iter.len() % (self.step + 1);
if self.first_take {
if rem == 0 { self.step } else { rem - 1 }
} else {
rem
}
}
}
#[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")]
impl<I> DoubleEndedIterator for StepBy<I> where I: DoubleEndedIterator + ExactSizeIterator {
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
self.iter.nth_back(self.next_back_index())
}
#[inline]
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
// `self.iter.nth_back(usize::MAX)` does the right thing here when `n`
// is out of bounds because the length of `self.iter` does not exceed
// `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is
// zero-indexed
let n = n
.saturating_mul(self.step + 1)
.saturating_add(self.next_back_index());
self.iter.nth_back(n)
}
}
// StepBy can only make the iterator shorter, so the len will still fit.
#[stable(feature = "iterator_step_by", since = "1.28.0")]
impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
@ -1158,6 +1191,45 @@ impl<I: Iterator> Iterator for Peekable<I> {
}
}
#[stable(feature = "double_ended_peek_iterator", since = "1.38.0")]
impl<I> DoubleEndedIterator for Peekable<I> where I: DoubleEndedIterator {
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
self.iter.next_back().or_else(|| self.peeked.take().and_then(|x| x))
}
#[inline]
fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R where
Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
{
match self.peeked.take() {
Some(None) => return Try::from_ok(init),
Some(Some(v)) => match self.iter.try_rfold(init, &mut f).into_result() {
Ok(acc) => f(acc, v),
Err(e) => {
self.peeked = Some(Some(v));
Try::from_error(e)
}
},
None => self.iter.try_rfold(init, f),
}
}
#[inline]
fn rfold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
where Fold: FnMut(Acc, Self::Item) -> Acc,
{
match self.peeked {
Some(None) => return init,
Some(Some(v)) => {
let acc = self.iter.rfold(init, &mut fold);
fold(acc, v)
}
None => self.iter.rfold(init, fold),
}
}
}
#[stable(feature = "rust1", since = "1.0.0")]
impl<I: ExactSizeIterator> ExactSizeIterator for Peekable<I> {}
@ -1613,6 +1685,51 @@ impl<I> Iterator for Take<I> where I: Iterator{
}
}
#[stable(feature = "double_ended_take_iterator", since = "1.38.0")]
impl<I> DoubleEndedIterator for Take<I> where I: DoubleEndedIterator + ExactSizeIterator {
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
if self.n == 0 {
None
} else {
let n = self.n;
self.n -= 1;
self.iter.nth_back(self.iter.len().saturating_sub(n))
}
}
#[inline]
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
let len = self.iter.len();
if self.n > n {
let m = len.saturating_sub(self.n) + n;
self.n -= n + 1;
self.iter.nth_back(m)
} else {
if len > 0 {
self.iter.nth_back(len - 1);
}
None
}
}
#[inline]
fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R where
Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try<Ok = Acc>
{
if self.n == 0 {
Try::from_ok(init)
} else {
let len = self.iter.len();
if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() {
Try::from_ok(init)
} else {
self.iter.try_rfold(init, fold)
}
}
}
}
#[stable(feature = "rust1", since = "1.0.0")]
impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {}

View File

@ -188,6 +188,19 @@ fn test_iterator_step_by() {
assert_eq!(it.next(), Some(6));
assert_eq!(it.next(), Some(9));
assert_eq!(it.next(), None);
let mut it = (0..3).step_by(1);
assert_eq!(it.next_back(), Some(2));
assert_eq!(it.next_back(), Some(1));
assert_eq!(it.next_back(), Some(0));
assert_eq!(it.next_back(), None);
let mut it = (0..11).step_by(3);
assert_eq!(it.next_back(), Some(9));
assert_eq!(it.next_back(), Some(6));
assert_eq!(it.next_back(), Some(3));
assert_eq!(it.next_back(), Some(0));
assert_eq!(it.next_back(), None);
}
#[test]
@ -252,6 +265,31 @@ fn test_iterator_step_by_nth_overflow() {
assert_eq!(it.0, (usize::MAX as Bigger) * 1);
}
#[test]
fn test_iterator_step_by_nth_back() {
let mut it = (0..16).step_by(5);
assert_eq!(it.nth_back(0), Some(15));
assert_eq!(it.nth_back(0), Some(10));
assert_eq!(it.nth_back(0), Some(5));
assert_eq!(it.nth_back(0), Some(0));
assert_eq!(it.nth_back(0), None);
let mut it = (0..16).step_by(5);
assert_eq!(it.next(), Some(0)); // to set `first_take` to `false`
assert_eq!(it.nth_back(0), Some(15));
assert_eq!(it.nth_back(0), Some(10));
assert_eq!(it.nth_back(0), Some(5));
assert_eq!(it.nth_back(0), None);
let it = || (0..18).step_by(5);
assert_eq!(it().nth_back(0), Some(15));
assert_eq!(it().nth_back(1), Some(10));
assert_eq!(it().nth_back(2), Some(5));
assert_eq!(it().nth_back(3), Some(0));
assert_eq!(it().nth_back(4), None);
assert_eq!(it().nth_back(42), None);
}
#[test]
#[should_panic]
fn test_iterator_step_by_zero() {
@ -465,8 +503,8 @@ fn test_iterator_filter_fold() {
#[test]
fn test_iterator_peekable() {
let xs = vec![0, 1, 2, 3, 4, 5];
let mut it = xs.iter().cloned().peekable();
let mut it = xs.iter().cloned().peekable();
assert_eq!(it.len(), 6);
assert_eq!(it.peek().unwrap(), &0);
assert_eq!(it.len(), 6);
@ -492,6 +530,33 @@ fn test_iterator_peekable() {
assert_eq!(it.len(), 0);
assert!(it.next().is_none());
assert_eq!(it.len(), 0);
let mut it = xs.iter().cloned().peekable();
assert_eq!(it.len(), 6);
assert_eq!(it.peek().unwrap(), &0);
assert_eq!(it.len(), 6);
assert_eq!(it.next_back().unwrap(), 5);
assert_eq!(it.len(), 5);
assert_eq!(it.next_back().unwrap(), 4);
assert_eq!(it.len(), 4);
assert_eq!(it.next_back().unwrap(), 3);
assert_eq!(it.len(), 3);
assert_eq!(it.peek().unwrap(), &0);
assert_eq!(it.len(), 3);
assert_eq!(it.peek().unwrap(), &0);
assert_eq!(it.len(), 3);
assert_eq!(it.next_back().unwrap(), 2);
assert_eq!(it.len(), 2);
assert_eq!(it.next_back().unwrap(), 1);
assert_eq!(it.len(), 1);
assert_eq!(it.peek().unwrap(), &0);
assert_eq!(it.len(), 1);
assert_eq!(it.next_back().unwrap(), 0);
assert_eq!(it.len(), 0);
assert!(it.peek().is_none());
assert_eq!(it.len(), 0);
assert!(it.next_back().is_none());
assert_eq!(it.len(), 0);
}
#[test]
@ -564,6 +629,18 @@ fn test_iterator_peekable_fold() {
assert_eq!(i, xs.len());
}
#[test]
fn test_iterator_peekable_rfold() {
let xs = [0, 1, 2, 3, 4, 5];
let mut it = xs.iter().peekable();
assert_eq!(it.peek(), Some(&&0));
let i = it.rfold(0, |i, &x| {
assert_eq!(x, xs[xs.len() - 1 - i]);
i + 1
});
assert_eq!(i, xs.len());
}
/// This is an iterator that follows the Iterator contract,
/// but it is not fused. After having returned None once, it will start
/// producing elements if .next() is called again.
@ -812,13 +889,25 @@ fn test_iterator_skip_fold() {
fn test_iterator_take() {
let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19];
let ys = [0, 1, 2, 3, 5];
let mut it = xs.iter().take(5);
let mut it = xs.iter().take(ys.len());
let mut i = 0;
assert_eq!(it.len(), 5);
assert_eq!(it.len(), ys.len());
while let Some(&x) = it.next() {
assert_eq!(x, ys[i]);
i += 1;
assert_eq!(it.len(), 5-i);
assert_eq!(it.len(), ys.len() - i);
}
assert_eq!(i, ys.len());
assert_eq!(it.len(), 0);
let mut it = xs.iter().take(ys.len());
let mut i = 0;
assert_eq!(it.len(), ys.len());
while let Some(&x) = it.next_back() {
i += 1;
assert_eq!(x, ys[ys.len() - i]);
assert_eq!(it.len(), ys.len() - i);
}
assert_eq!(i, ys.len());
assert_eq!(it.len(), 0);
@ -848,19 +937,51 @@ fn test_iterator_take_nth() {
}
}
#[test]
fn test_iterator_take_nth_back() {
let xs = [0, 1, 2, 4, 5];
let mut it = xs.iter();
{
let mut take = it.by_ref().take(3);
let mut i = 0;
while let Some(&x) = take.nth_back(0) {
i += 1;
assert_eq!(x, 3 - i);
}
}
assert_eq!(it.nth_back(0), None);
let xs = [0, 1, 2, 3, 4];
let mut it = xs.iter().take(7);
assert_eq!(it.nth_back(1), Some(&3));
assert_eq!(it.nth_back(1), Some(&1));
assert_eq!(it.nth_back(1), None);
}
#[test]
fn test_iterator_take_short() {
let xs = [0, 1, 2, 3];
let ys = [0, 1, 2, 3];
let mut it = xs.iter().take(5);
let mut i = 0;
assert_eq!(it.len(), 4);
assert_eq!(it.len(), xs.len());
while let Some(&x) = it.next() {
assert_eq!(x, ys[i]);
assert_eq!(x, xs[i]);
i += 1;
assert_eq!(it.len(), 4-i);
assert_eq!(it.len(), xs.len() - i);
}
assert_eq!(i, ys.len());
assert_eq!(i, xs.len());
assert_eq!(it.len(), 0);
let mut it = xs.iter().take(5);
let mut i = 0;
assert_eq!(it.len(), xs.len());
while let Some(&x) = it.next_back() {
i += 1;
assert_eq!(x, xs[xs.len() - i]);
assert_eq!(it.len(), xs.len() - i);
}
assert_eq!(i, xs.len());
assert_eq!(it.len(), 0);
}
@ -2241,17 +2362,50 @@ fn test_enumerate_try_folds() {
}
#[test]
fn test_peek_try_fold() {
fn test_peek_try_folds() {
let f = &|acc, x| i32::checked_add(2*acc, x);
assert_eq!((1..20).peekable().try_fold(7, f), (1..20).try_fold(7, f));
assert_eq!((1..20).peekable().try_rfold(7, f), (1..20).try_rfold(7, f));
let mut iter = (1..20).peekable();
assert_eq!(iter.peek(), Some(&1));
assert_eq!(iter.try_fold(7, f), (1..20).try_fold(7, f));
let mut iter = (1..20).peekable();
assert_eq!(iter.peek(), Some(&1));
assert_eq!(iter.try_rfold(7, f), (1..20).try_rfold(7, f));
let mut iter = [100, 20, 30, 40, 50, 60, 70].iter().cloned().peekable();
assert_eq!(iter.peek(), Some(&100));
assert_eq!(iter.try_fold(0, i8::checked_add), None);
assert_eq!(iter.peek(), Some(&40));
let mut iter = [100, 20, 30, 40, 50, 60, 70].iter().cloned().peekable();
assert_eq!(iter.peek(), Some(&100));
assert_eq!(iter.try_rfold(0, i8::checked_add), None);
assert_eq!(iter.peek(), Some(&100));
assert_eq!(iter.next_back(), Some(50));
let mut iter = (2..5).peekable();
assert_eq!(iter.peek(), Some(&2));
assert_eq!(iter.try_for_each(Err), Err(2));
assert_eq!(iter.peek(), Some(&3));
assert_eq!(iter.try_for_each(Err), Err(3));
assert_eq!(iter.peek(), Some(&4));
assert_eq!(iter.try_for_each(Err), Err(4));
assert_eq!(iter.peek(), None);
assert_eq!(iter.try_for_each(Err), Ok(()));
let mut iter = (2..5).peekable();
assert_eq!(iter.peek(), Some(&2));
assert_eq!(iter.try_rfold((), |(), x| Err(x)), Err(4));
assert_eq!(iter.peek(), Some(&2));
assert_eq!(iter.try_rfold((), |(), x| Err(x)), Err(3));
assert_eq!(iter.peek(), Some(&2));
assert_eq!(iter.try_rfold((), |(), x| Err(x)), Err(2));
assert_eq!(iter.peek(), None);
assert_eq!(iter.try_rfold((), |(), x| Err(x)), Ok(()));
}
#[test]
@ -2300,13 +2454,25 @@ fn test_skip_try_folds() {
fn test_take_try_folds() {
let f = &|acc, x| i32::checked_add(2*acc, x);
assert_eq!((10..30).take(10).try_fold(7, f), (10..20).try_fold(7, f));
//assert_eq!((10..30).take(10).try_rfold(7, f), (10..20).try_rfold(7, f));
assert_eq!((10..30).take(10).try_rfold(7, f), (10..20).try_rfold(7, f));
let mut iter = (10..30).take(20);
assert_eq!(iter.try_fold(0, i8::checked_add), None);
assert_eq!(iter.next(), Some(20));
//assert_eq!(iter.try_rfold(0, i8::checked_add), None);
//assert_eq!(iter.next_back(), Some(24));
assert_eq!(iter.try_rfold(0, i8::checked_add), None);
assert_eq!(iter.next_back(), Some(24));
let mut iter = (2..20).take(3);
assert_eq!(iter.try_for_each(Err), Err(2));
assert_eq!(iter.try_for_each(Err), Err(3));
assert_eq!(iter.try_for_each(Err), Err(4));
assert_eq!(iter.try_for_each(Err), Ok(()));
let mut iter = (2..20).take(3).rev();
assert_eq!(iter.try_for_each(Err), Err(4));
assert_eq!(iter.try_for_each(Err), Err(3));
assert_eq!(iter.try_for_each(Err), Err(2));
assert_eq!(iter.try_for_each(Err), Ok(()));
}
#[test]