diff options
| -rw-r--r-- | library/core/src/iter/adapters/chain.rs | 76 | ||||
| -rw-r--r-- | library/core/tests/iter.rs | 103 |
2 files changed, 167 insertions, 12 deletions
diff --git a/library/core/src/iter/adapters/chain.rs b/library/core/src/iter/adapters/chain.rs index ac27ec19b36..38fb74372db 100644 --- a/library/core/src/iter/adapters/chain.rs +++ b/library/core/src/iter/adapters/chain.rs @@ -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)) } diff --git a/library/core/tests/iter.rs b/library/core/tests/iter.rs index b15d6d1b1f6..75ca897cadc 100644 --- a/library/core/tests/iter.rs +++ b/library/core/tests/iter.rs @@ -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] = []; @@ -143,6 +180,72 @@ fn test_iterator_chain() { } #[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]; let ys = [30, 40, 50, 60]; |