replace advance_by returning usize with Result<(), NonZeroUsize>

Author: the8472

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) -> usize {
-        if self.inner.len < n {
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
+        let rem = if self.inner.len < n {
             let len = self.inner.len;
             self.inner.clear();
             len - n
         } else {
             self.inner.drain(..n);
             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) -> usize {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         let len = self.inner.len;
-        if len >= n {
+        let rem = if len >= n {
             self.inner.truncate(len - n);
             0
         } else {
             self.inner.clear();
             n - len
-        }
+        };
+        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) -> usize {
+    fn advance_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         let remaining = self.i1.advance_by(n);
-        if remaining == 0 {
-            return 0;
+        match remaining {
+            Ok(()) => return Ok(()),
+            Err(n) => {
+                mem::swap(&mut self.i1, &mut self.i2);
+                self.i1.advance_by(n.get())
+            }
         }
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i1.advance_by(remaining)
     }
 
     #[inline]
@@ -125,13 +128,14 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
         }
     }
 
-    fn advance_back_by(&mut self, n: usize) -> usize {
-        let remaining = self.i2.advance_back_by(n);
-        if remaining == 0 {
-            return 0;
+    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())
+            }
         }
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i2.advance_back_by(remaining)
     }
 
     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) -> usize {
-        let remaining = self.i1.advance_by(n);
-        if remaining == 0 {
-            return 0;
+    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())
+            }
         }
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i1.advance_by(remaining)
     }
 
     #[inline]
@@ -117,13 +119,14 @@ impl<'a, T> DoubleEndedIterator for IterMut<'a, T> {
         }
     }
 
-    fn advance_back_by(&mut self, n: usize) -> usize {
-        let remaining = self.i2.advance_back_by(n);
-        if remaining == 0 {
-            return 0;
+    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())
+            }
         }
-        mem::swap(&mut self.i1, &mut self.i2);
-        self.i2.advance_back_by(remaining)
     }
 
     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) -> 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,7 +228,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         unsafe {
             ptr::drop_in_place(to_drop);
         }
-        n - step_size
+        NonZeroUsize::new(n - step_size).map_or(Ok(()), Err)
     }
 
     #[inline]
@@ -310,7 +311,7 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
     }
 
     #[inline]
-    fn advance_back_by(&mut self, n: usize) -> 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()
@@ -324,7 +325,7 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
         unsafe {
             ptr::drop_in_place(to_drop);
         }
-        n - step_size
+        NonZeroUsize::new(n - step_size).map_or(Ok(()), Err)
     }
 }
 

library/alloc/tests/vec.rs

@@ -1,6 +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;
@@ -1064,20 +1065,20 @@ fn test_into_iter_leak() {
 #[test]
 fn test_into_iter_advance_by() {
     let mut i = vec![1, 2, 3, 4, 5].into_iter();
-    assert_eq!(i.advance_by(0), 0);
-    assert_eq!(i.advance_back_by(0), 0);
+    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]);
 
-    assert_eq!(i.advance_by(1), 0);
-    assert_eq!(i.advance_back_by(1), 0);
+    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), usize::MAX - 3);
+    assert_eq!(i.advance_back_by(usize::MAX), Err(NonZeroUsize::new(usize::MAX - 3).unwrap()));
 
-    assert_eq!(i.advance_by(usize::MAX), usize::MAX);
+    assert_eq!(i.advance_by(usize::MAX), Err(NonZeroUsize::new(usize::MAX).unwrap()));
 
-    assert_eq!(i.advance_by(0), 0);
-    assert_eq!(i.advance_back_by(0), 0);
+    assert_eq!(i.advance_by(0), Ok(()));
+    assert_eq!(i.advance_back_by(0), Ok(()));
 
     assert_eq!(i.len(), 0);
 }
@@ -1125,7 +1126,7 @@ fn test_into_iter_zst() {
     for _ in vec![C; 5].into_iter().rev() {}
 
     let mut it = vec![C, C].into_iter();
-    assert_eq!(it.advance_by(1), 0);
+    assert_eq!(it.advance_by(1), Ok(()));
     drop(it);
 
     let mut it = vec![C, C].into_iter();

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,7 +285,7 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
         self.next_back()
     }
 
-    fn advance_by(&mut self, n: usize) -> usize {
+    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);
@@ -296,7 +297,7 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
             ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
         }
 
-        remaining
+        NonZeroUsize::new(remaining).map_or(Ok(()), Err)
     }
 }
 
@@ -333,7 +334,7 @@ impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
         })
     }
 
-    fn advance_back_by(&mut self, n: usize) -> usize {
+    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);
@@ -345,7 +346,7 @@ impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
             ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
         }
 
-        remaining
+        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) -> 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) -> usize {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> {
         I::advance_back_by(self.0, n)
     }
 

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

@@ -1,4 +1,5 @@
 use crate::iter::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen};
+use crate::num::NonZeroUsize;
 use crate::ops::Try;
 
 /// An iterator that links two iterators together, in a chain.
@@ -95,32 +96,32 @@ where
     }
 
     #[inline]
-    fn advance_by(&mut self, mut n: usize) -> usize {
+    fn advance_by(&mut self, mut n: usize) -> Result<(), NonZeroUsize> {
         if let Some(ref mut a) = self.a {
-            n = a.advance_by(n);
-            if n == 0 {
-                return n;
-            }
+            n = match a.advance_by(n) {
+                Ok(()) => return Ok(()),
+                Err(k) => k.get(),
+            };
             self.a = None;
         }
 
         if let Some(ref mut b) = self.b {
-            n = b.advance_by(n);
+            return b.advance_by(n);
             // we don't fuse the second iterator
         }
 
-        n
+        NonZeroUsize::new(n).map_or(Ok(()), Err)
     }
 
     #[inline]
     fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
         if let Some(ref mut a) = self.a {
             n = match a.advance_by(n) {
-                0 => match a.next() {
+                Ok(()) => match a.next() {
                     None => 0,
                     x => return x,
                 },
-                k => k,
+                Err(k) => k.get(),
             };
 
             self.a = None;
@@ -181,32 +182,32 @@ where
     }
 
     #[inline]
-    fn advance_back_by(&mut self, mut n: usize) -> usize {
+    fn advance_back_by(&mut self, mut n: usize) -> Result<(), NonZeroUsize> {
         if let Some(ref mut b) = self.b {
-            n = b.advance_back_by(n);
-            if n == 0 {
-                return n;
-            }
+            n = match b.advance_back_by(n) {
+                Ok(()) => return Ok(()),
+                Err(k) => k.get(),
+            };
             self.b = None;
         }
 
         if let Some(ref mut a) = self.a {
-            n = a.advance_back_by(n);
+            return a.advance_back_by(n);
             // we don't fuse the second iterator
         }
 
-        n
+        NonZeroUsize::new(n).map_or(Ok(()), Err)
     }
 
     #[inline]
     fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
         if let Some(ref mut b) = self.b {
             n = match b.advance_back_by(n) {
-                0 => match b.next_back() {
+                Ok(()) => match b.next_back() {
                     None => 0,
                     x => return x,
                 },
-                k => k,
+                Err(k) => k.get(),
             };
 
             self.b = None;