Change advance(_back)_by to return the remainder instead of the number of processed elements

library/alloc/src/collections/vec_deque/into_iter.rs

@@ -1,4 +1,5 @@
 use core::iter::{FusedIterator, TrustedLen};
+use core::num::NonZeroUsize;
 use core::{array, fmt, mem::MaybeUninit, ops::Try, ptr};
 
 use crate::alloc::{Allocator, Global};
@@ -54,15 +55,16 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
     }
 
     #[inline]
-    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
-        if self.inner.len < n {
-            let len = self.inner.len;
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
+        let len = self.inner.len;
+        let rem = if len < n {
             self.inner.clear();
-            Err(len)
+            n - len
         } else {
             self.inner.drain(..n);
-            Ok(())
-        }
+            0
+        };
+        NonZeroUsize::new(rem).map_or(Ok(()), Err)
     }
 
     #[inline]
@@ -182,15 +184,16 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
     }
 
     #[inline]
-    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         let len = self.inner.len;
-        if len >= n {
-            self.inner.truncate(len - n);
-            Ok(())
-        } else {
+        let rem = if len < n {
             self.inner.clear();
-            Err(len)
-        }
+            n - len
+        } else {
+            self.inner.truncate(len - n);
+            0
+        };
+        NonZeroUsize::new(rem).map_or(Ok(()), Err)
     }
 
     fn try_rfold<B, F, R>(&mut self, mut init: B, mut f: F) -> R

library/alloc/src/collections/vec_deque/iter.rs

@@ -1,4 +1,5 @@
 use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce};
+use core::num::NonZeroUsize;
 use core::ops::Try;
 use core::{fmt, mem, slice};
 
@@ -55,13 +56,15 @@ impl<'a, T> Iterator for Iter<'a, T> {
         }
     }
 
-    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
-        let m = match self.i1.advance_by(n) {
-            Ok(_) => return Ok(()),
-            Err(m) => m,
-        };
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i1.advance_by(n - m).map_err(|o| o + m)
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
+        let remaining = self.i1.advance_by(n);
+        match remaining {
+            Ok(()) => return Ok(()),
+            Err(n) => {
+                mem::swap(&mut self.i1, &mut self.i2);
+                self.i1.advance_by(n.get())
+            }
+        }
     }
 
     #[inline]
@@ -125,14 +128,14 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
         }
     }
 
-    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
-        let m = match self.i2.advance_back_by(n) {
-            Ok(_) => return Ok(()),
-            Err(m) => m,
-        };
-
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i2.advance_back_by(n - m).map_err(|o| m + o)
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
+        match self.i2.advance_back_by(n) {
+            Ok(()) => return Ok(()),
+            Err(n) => {
+                mem::swap(&mut self.i1, &mut self.i2);
+                self.i2.advance_back_by(n.get())
+            }
+        }
     }
 
     fn rfold<Acc, F>(self, accum: Acc, mut f: F) -> Acc

library/alloc/src/collections/vec_deque/iter_mut.rs

@@ -1,4 +1,5 @@
 use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce};
+use core::num::NonZeroUsize;
 use core::ops::Try;
 use core::{fmt, mem, slice};
 
@@ -47,13 +48,14 @@ impl<'a, T> Iterator for IterMut<'a, T> {
         }
     }
 
-    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
-        let m = match self.i1.advance_by(n) {
-            Ok(_) => return Ok(()),
-            Err(m) => m,
-        };
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i1.advance_by(n - m).map_err(|o| o + m)
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
+        match self.i1.advance_by(n) {
+            Ok(()) => return Ok(()),
+            Err(remaining) => {
+                mem::swap(&mut self.i1, &mut self.i2);
+                self.i1.advance_by(remaining.get())
+            }
+        }
     }
 
     #[inline]
@@ -117,14 +119,14 @@ impl<'a, T> DoubleEndedIterator for IterMut<'a, T> {
         }
     }
 
-    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
-        let m = match self.i2.advance_back_by(n) {
-            Ok(_) => return Ok(()),
-            Err(m) => m,
-        };
-
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i2.advance_back_by(n - m).map_err(|o| m + o)
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
+        match self.i2.advance_back_by(n) {
+            Ok(()) => return Ok(()),
+            Err(remaining) => {
+                mem::swap(&mut self.i1, &mut self.i2);
+                self.i2.advance_back_by(remaining.get())
+            }
+        }
     }
 
     fn rfold<Acc, F>(self, accum: Acc, mut f: F) -> Acc

library/alloc/src/vec/into_iter.rs

@@ -11,6 +11,7 @@ use core::iter::{
 };
 use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties};
+use core::num::NonZeroUsize;
 #[cfg(not(no_global_oom_handling))]
 use core::ops::Deref;
 use core::ptr::{self, NonNull};
@@ -213,7 +214,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
     }
 
     #[inline]
-    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         let step_size = self.len().min(n);
         let to_drop = ptr::slice_from_raw_parts_mut(self.ptr as *mut T, step_size);
         if T::IS_ZST {
@@ -227,10 +228,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         unsafe {
             ptr::drop_in_place(to_drop);
         }
-        if step_size < n {
-            return Err(step_size);
-        }
-        Ok(())
+        NonZeroUsize::new(n - step_size).map_or(Ok(()), Err)
     }
 
     #[inline]
@@ -313,7 +311,7 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
     }
 
     #[inline]
-    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         let step_size = self.len().min(n);
         if T::IS_ZST {
             // SAFETY: same as for advance_by()
@@ -327,10 +325,7 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
         unsafe {
             ptr::drop_in_place(to_drop);
         }
-        if step_size < n {
-            return Err(step_size);
-        }
-        Ok(())
+        NonZeroUsize::new(n - step_size).map_or(Ok(()), Err)
     }
 }
 

library/alloc/tests/vec.rs

@@ -1,5 +1,7 @@
 use core::alloc::{Allocator, Layout};
+use core::assert_eq;
 use core::iter::IntoIterator;
+use core::num::NonZeroUsize;
 use core::ptr::NonNull;
 use std::alloc::System;
 use std::assert_matches::assert_matches;
@@ -1062,21 +1064,21 @@ fn test_into_iter_leak() {
 
 #[test]
 fn test_into_iter_advance_by() {
-    let mut i = [1, 2, 3, 4, 5].into_iter();
-    i.advance_by(0).unwrap();
-    i.advance_back_by(0).unwrap();
+    let mut i = vec![1, 2, 3, 4, 5].into_iter();
+    assert_eq!(i.advance_by(0), Ok(()));
+    assert_eq!(i.advance_back_by(0), Ok(()));
     assert_eq!(i.as_slice(), [1, 2, 3, 4, 5]);
 
-    i.advance_by(1).unwrap();
-    i.advance_back_by(1).unwrap();
+    assert_eq!(i.advance_by(1), Ok(()));
+    assert_eq!(i.advance_back_by(1), Ok(()));
     assert_eq!(i.as_slice(), [2, 3, 4]);
 
-    assert_eq!(i.advance_back_by(usize::MAX), Err(3));
+    assert_eq!(i.advance_back_by(usize::MAX), Err(NonZeroUsize::new(usize::MAX - 3).unwrap()));
 
-    assert_eq!(i.advance_by(usize::MAX), Err(0));
+    assert_eq!(i.advance_by(usize::MAX), Err(NonZeroUsize::new(usize::MAX).unwrap()));
 
-    i.advance_by(0).unwrap();
-    i.advance_back_by(0).unwrap();
+    assert_eq!(i.advance_by(0), Ok(()));
+    assert_eq!(i.advance_back_by(0), Ok(()));
 
     assert_eq!(i.len(), 0);
 }
@@ -1124,7 +1126,7 @@ fn test_into_iter_zst() {
     for _ in vec![C; 5].into_iter().rev() {}
 
     let mut it = vec![C, C].into_iter();
-    it.advance_by(1).unwrap();
+    assert_eq!(it.advance_by(1), Ok(()));
     drop(it);
 
     let mut it = vec![C, C].into_iter();

library/alloc/tests/vec_deque.rs

@@ -1,3 +1,4 @@
+use core::num::NonZeroUsize;
 use std::assert_matches::assert_matches;
 use std::collections::TryReserveErrorKind::*;
 use std::collections::{vec_deque::Drain, VecDeque};
@@ -426,6 +427,28 @@ fn test_into_iter() {
         assert_eq!(it.next(), Some(7));
         assert_eq!(it.size_hint(), (5, Some(5)));
     }
+
+    // advance_by
+    {
+        let mut d = VecDeque::new();
+        for i in 0..=4 {
+            d.push_back(i);
+        }
+        for i in 6..=8 {
+            d.push_front(i);
+        }
+
+        let mut it = d.into_iter();
+        assert_eq!(it.advance_by(1), Ok(()));
+        assert_eq!(it.next(), Some(7));
+        assert_eq!(it.advance_back_by(1), Ok(()));
+        assert_eq!(it.next_back(), Some(3));
+
+        let mut it = VecDeque::from(vec![1, 2, 3, 4, 5]).into_iter();
+        assert_eq!(it.advance_by(10), Err(NonZeroUsize::new(5).unwrap()));
+        let mut it = VecDeque::from(vec![1, 2, 3, 4, 5]).into_iter();
+        assert_eq!(it.advance_back_by(10), Err(NonZeroUsize::new(5).unwrap()));
+    }
 }
 
 #[test]

library/core/src/array/iter.rs

@@ -1,5 +1,6 @@
 //! Defines the `IntoIter` owned iterator for arrays.
 
+use crate::num::NonZeroUsize;
 use crate::{
     fmt,
     iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
@@ -284,20 +285,19 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
         self.next_back()
     }
 
-    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
-        let original_len = self.len();
-
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         // This also moves the start, which marks them as conceptually "dropped",
         // so if anything goes bad then our drop impl won't double-free them.
         let range_to_drop = self.alive.take_prefix(n);
+        let remaining = n - range_to_drop.len();
 
         // SAFETY: These elements are currently initialized, so it's fine to drop them.
         unsafe {
             let slice = self.data.get_unchecked_mut(range_to_drop);
             ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
         }
 
-        if n > original_len { Err(original_len) } else { Ok(()) }
+        NonZeroUsize::new(remaining).map_or(Ok(()), Err)
     }
 }
 
@@ -334,20 +334,19 @@ impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
         })
     }
 
-    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
-        let original_len = self.len();
-
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         // This also moves the end, which marks them as conceptually "dropped",
         // so if anything goes bad then our drop impl won't double-free them.
         let range_to_drop = self.alive.take_suffix(n);
+        let remaining = n - range_to_drop.len();
 
         // SAFETY: These elements are currently initialized, so it's fine to drop them.
         unsafe {
             let slice = self.data.get_unchecked_mut(range_to_drop);
             ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
         }
 
-        if n > original_len { Err(original_len) } else { Ok(()) }
+        NonZeroUsize::new(remaining).map_or(Ok(()), Err)
     }
 }
 

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

@@ -1,3 +1,4 @@
+use crate::num::NonZeroUsize;
 use crate::ops::{NeverShortCircuit, Try};
 
 /// Like `Iterator::by_ref`, but requiring `Sized` so it can forward generics.
@@ -26,7 +27,7 @@ impl<I: Iterator> Iterator for ByRefSized<'_, I> {
     }
 
     #[inline]
-    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         I::advance_by(self.0, n)
     }
 
@@ -62,7 +63,7 @@ impl<I: DoubleEndedIterator> DoubleEndedIterator for ByRefSized<'_, I> {
     }
 
     #[inline]
-    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         I::advance_back_by(self.0, n)
     }