diff --git a/library/core/src/iter/adapters/chain.rs b/library/core/src/iter/adapters/chain.rs index 13c6a75d58b..dadd7bfb0f7 100644 --- a/library/core/src/iter/adapters/chain.rs +++ b/library/core/src/iter/adapters/chain.rs @@ -115,16 +115,42 @@ where } #[inline] - fn nth(&mut self, mut n: usize) -> Option { + 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 { + 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 { + 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 { + 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)) }