diff --git a/library/core/src/slice/iter/macros.rs b/library/core/src/slice/iter/macros.rs index 457b2a3605e8..be7d09e32603 100644 --- a/library/core/src/slice/iter/macros.rs +++ b/library/core/src/slice/iter/macros.rs @@ -163,26 +163,11 @@ macro_rules! iterator { } #[inline] - fn nth(&mut self, n: usize) -> Option<$elem> { - if n >= len!(self) { - // This iterator is now empty. - if mem::size_of::() == 0 { - // We have to do it this way as `ptr` may never be 0, but `end` - // could be (due to wrapping). - self.end = self.ptr.as_ptr(); - } else { - // SAFETY: end can't be 0 if T isn't ZST because ptr isn't 0 and end >= ptr - unsafe { - self.ptr = NonNull::new_unchecked(self.end as *mut T); - } - } - return None; - } - // SAFETY: We are in bounds. `post_inc_start` does the right thing even for ZSTs. - unsafe { - self.post_inc_start(n as isize); - Some(next_unchecked!(self)) - } + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + let advance = cmp::min(n, len!(self)); + // SAFETY: By construction, `advance` does not exceed `self.len()`. + unsafe { self.post_inc_start(advance as isize) }; + if advance == n { Ok(()) } else { Err(advance) } } #[inline] @@ -359,17 +344,11 @@ macro_rules! iterator { } #[inline] - fn nth_back(&mut self, n: usize) -> Option<$elem> { - if n >= len!(self) { - // This iterator is now empty. - self.end = self.ptr.as_ptr(); - return None; - } - // SAFETY: We are in bounds. `pre_dec_end` does the right thing even for ZSTs. - unsafe { - self.pre_dec_end(n as isize); - Some(next_back_unchecked!(self)) - } + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + let advance = cmp::min(n, len!(self)); + // SAFETY: By construction, `advance` does not exceed `self.len()`. + unsafe { self.pre_dec_end(advance as isize) }; + if advance == n { Ok(()) } else { Err(advance) } } } diff --git a/library/core/tests/slice.rs b/library/core/tests/slice.rs index 5ef30b1a8898..f6b1d09ddbff 100644 --- a/library/core/tests/slice.rs +++ b/library/core/tests/slice.rs @@ -122,6 +122,48 @@ fn test_partition_point() { assert_eq!(b.partition_point(|&x| x < 8), 5); } +#[test] +fn test_iterator_advance_by() { + let v = &[0, 1, 2, 3, 4]; + + for i in 0..=v.len() { + let mut iter = v.iter(); + iter.advance_by(i).unwrap(); + assert_eq!(iter.as_slice(), &v[i..]); + } + + let mut iter = v.iter(); + assert_eq!(iter.advance_by(v.len() + 1), Err(v.len())); + assert_eq!(iter.as_slice(), &[]); + + let mut iter = v.iter(); + iter.advance_by(3).unwrap(); + assert_eq!(iter.as_slice(), &v[3..]); + iter.advance_by(2).unwrap(); + assert_eq!(iter.as_slice(), &[]); +} + +#[test] +fn test_iterator_advance_back_by() { + let v = &[0, 1, 2, 3, 4]; + + for i in 0..=v.len() { + let mut iter = v.iter(); + iter.advance_back_by(i).unwrap(); + assert_eq!(iter.as_slice(), &v[..v.len() - i]); + } + + let mut iter = v.iter(); + assert_eq!(iter.advance_back_by(v.len() + 1), Err(v.len())); + assert_eq!(iter.as_slice(), &[]); + + let mut iter = v.iter(); + iter.advance_back_by(3).unwrap(); + assert_eq!(iter.as_slice(), &v[..v.len() - 3]); + iter.advance_back_by(2).unwrap(); + assert_eq!(iter.as_slice(), &[]); +} + #[test] fn test_iterator_nth() { let v: &[_] = &[0, 1, 2, 3, 4];