Implement advance_by, advance_back_by for iter::Chain

library/core/src/iter/adapters/chain.rs

@@ -115,16 +115,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))
  }
 
@@ -191,16 +217,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))
  }
 

library/core/tests/iter.rs

@@ -142,6 +142,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 = xs.iter().chain(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 = xs.iter().chain(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 = xs.iter().chain(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 = xs.iter().chain(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];