Implement advance_by, advance_back_by for slice::{Iter, IterMut}

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::<T>() == 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) }
             }
         }
 

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];