Auto merge of #77594 - timvermeulen:chain_advance_by, r=scottmcm

Implement advance_by, advance_back_by for iter::Chain

Part of #77404.

This PR does two things:
- implement `Chain::advance[_back]_by` in terms of `advance[_back]_by` on `self.a` and `advance[_back]_by` on `self.b`
- change `Chain::nth[_back]` to use `advance[_back]_by` on `self.a` and `nth[_back]` on `self.b`

This ensures that `Chain::nth` can take advantage of an efficient `nth` implementation on the second iterator, in case it doesn't implement `advance_by`.

cc `@scottmcm` in case you want to review this
This commit is contained in:
bors 2020-10-06 10:17:48 +00:00
commit 5849a7eca9
2 changed files with 167 additions and 12 deletions

View File

@ -126,16 +126,42 @@ where
}
#[inline]
fn nth(&mut self, mut n: usize) -> Option<A::Item> {
fn advance_by(&mut self, n: usize) -> Result<(), usize> {
let mut rem = n;
if let Some(ref mut a) = self.a {
while let Some(x) = a.next() {
if n == 0 {
return Some(x);
}
n -= 1;
match a.advance_by(rem) {
Ok(()) => return Ok(()),
Err(k) => rem -= k,
}
self.a = None;
}
if let Some(ref mut b) = self.b {
match b.advance_by(rem) {
Ok(()) => return Ok(()),
Err(k) => rem -= k,
}
// we don't fuse the second iterator
}
if rem == 0 { Ok(()) } else { Err(n - rem) }
}
#[inline]
fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
if let Some(ref mut a) = self.a {
match a.advance_by(n) {
Ok(()) => match a.next() {
None => n = 0,
x => return x,
},
Err(k) => n -= k,
}
self.a = None;
}
maybe!(self.b.nth(n))
}
@ -202,16 +228,42 @@ where
}
#[inline]
fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
let mut rem = n;
if let Some(ref mut b) = self.b {
while let Some(x) = b.next_back() {
if n == 0 {
return Some(x);
}
n -= 1;
match b.advance_back_by(rem) {
Ok(()) => return Ok(()),
Err(k) => rem -= k,
}
self.b = None;
}
if let Some(ref mut a) = self.a {
match a.advance_back_by(rem) {
Ok(()) => return Ok(()),
Err(k) => rem -= k,
}
// we don't fuse the second iterator
}
if rem == 0 { Ok(()) } else { Err(n - rem) }
}
#[inline]
fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
if let Some(ref mut b) = self.b {
match b.advance_back_by(n) {
Ok(()) => match b.next_back() {
None => n = 0,
x => return x,
},
Err(k) => n -= k,
}
self.b = None;
}
maybe!(self.a.nth_back(n))
}

View File

@ -4,6 +4,43 @@ use core::cell::Cell;
use core::convert::TryFrom;
use core::iter::*;
/// An iterator wrapper that panics whenever `next` or `next_back` is called
/// after `None` has been returned.
struct Unfuse<I> {
iter: I,
exhausted: bool,
}
fn unfuse<I: IntoIterator>(iter: I) -> Unfuse<I::IntoIter> {
Unfuse { iter: iter.into_iter(), exhausted: false }
}
impl<I> Iterator for Unfuse<I>
where
I: Iterator,
{
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
assert!(!self.exhausted);
let next = self.iter.next();
self.exhausted = next.is_none();
next
}
}
impl<I> DoubleEndedIterator for Unfuse<I>
where
I: DoubleEndedIterator,
{
fn next_back(&mut self) -> Option<Self::Item> {
assert!(!self.exhausted);
let next = self.iter.next_back();
self.exhausted = next.is_none();
next
}
}
#[test]
fn test_lt() {
let empty: [isize; 0] = [];
@ -142,6 +179,72 @@ fn test_iterator_chain() {
assert_eq!(i, expected.len());
}
#[test]
fn test_iterator_chain_advance_by() {
fn test_chain(xs: &[i32], ys: &[i32]) {
let len = xs.len() + ys.len();
for i in 0..xs.len() {
let mut iter = unfuse(xs).chain(unfuse(ys));
iter.advance_by(i).unwrap();
assert_eq!(iter.next(), Some(&xs[i]));
assert_eq!(iter.advance_by(100), Err(len - i - 1));
}
for i in 0..ys.len() {
let mut iter = unfuse(xs).chain(unfuse(ys));
iter.advance_by(xs.len() + i).unwrap();
assert_eq!(iter.next(), Some(&ys[i]));
assert_eq!(iter.advance_by(100), Err(ys.len() - i - 1));
}
let mut iter = xs.iter().chain(ys);
iter.advance_by(len).unwrap();
assert_eq!(iter.next(), None);
let mut iter = xs.iter().chain(ys);
assert_eq!(iter.advance_by(len + 1), Err(len));
}
test_chain(&[], &[]);
test_chain(&[], &[0, 1, 2, 3, 4, 5]);
test_chain(&[0, 1, 2, 3, 4, 5], &[]);
test_chain(&[0, 1, 2, 3, 4, 5], &[30, 40, 50, 60]);
}
#[test]
fn test_iterator_chain_advance_back_by() {
fn test_chain(xs: &[i32], ys: &[i32]) {
let len = xs.len() + ys.len();
for i in 0..ys.len() {
let mut iter = unfuse(xs).chain(unfuse(ys));
iter.advance_back_by(i).unwrap();
assert_eq!(iter.next_back(), Some(&ys[ys.len() - i - 1]));
assert_eq!(iter.advance_back_by(100), Err(len - i - 1));
}
for i in 0..xs.len() {
let mut iter = unfuse(xs).chain(unfuse(ys));
iter.advance_back_by(ys.len() + i).unwrap();
assert_eq!(iter.next_back(), Some(&xs[xs.len() - i - 1]));
assert_eq!(iter.advance_back_by(100), Err(xs.len() - i - 1));
}
let mut iter = xs.iter().chain(ys);
iter.advance_back_by(len).unwrap();
assert_eq!(iter.next_back(), None);
let mut iter = xs.iter().chain(ys);
assert_eq!(iter.advance_back_by(len + 1), Err(len));
}
test_chain(&[], &[]);
test_chain(&[], &[0, 1, 2, 3, 4, 5]);
test_chain(&[0, 1, 2, 3, 4, 5], &[]);
test_chain(&[0, 1, 2, 3, 4, 5], &[30, 40, 50, 60]);
}
#[test]
fn test_iterator_chain_nth() {
let xs = [0, 1, 2, 3, 4, 5];