Rollup merge of #92284 - the8472:simplify-advance-by, r=scottmcm

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

When advance_by can't advance the iterator by the number of requested elements it now returns the amount by which it couldn't be advanced instead of the amount by which it did.

This simplifies adapters like chain, flatten or cycle because the remainder doesn't have to be calculated as the difference between requested steps and completed steps anymore.

Additionally switching from `Result<(), usize>` to `Result<(), NonZeroUsize>` reduces the size of the result and makes converting from/to a usize representing the number of remaining steps cheap.

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)
     }