Rollup merge of #87091 - the8472:more-advance-by-impls, r=joshtriplett

implement advance_(back_)_by on more iterators

Add more efficient, non-default implementations for `feature(iter_advance_by)` (#77404) on more iterators and adapters.

This PR only contains implementations where skipping over items doesn't elide any observable side-effects such as user-provided closures or `clone()` functions. I'll put those in a separate PR.

4 files changed, 66 insertions, 0 deletions

diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index 2aed9d03bc0..ca41ce975e4 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -111,6 +111,7 @@
 // that the feature-gate isn't enabled. Ideally, it wouldn't check for the feature gate for docs
 // from other crates, but since this can only appear for lang items, it doesn't seem worth fixing.
 #![feature(intra_doc_pointers)]
+#![feature(iter_advance_by)]
 #![feature(iter_zip)]
 #![feature(lang_items)]
 #![feature(layout_for_ptr)]
diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs
index 4cb0a4b10bd..8a2d254a834 100644
--- a/library/alloc/src/vec/into_iter.rs
+++ b/library/alloc/src/vec/into_iter.rs
@@ -162,6 +162,29 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
     }
 
     #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        let step_size = self.len().min(n);
+        let to_drop = ptr::slice_from_raw_parts_mut(self.ptr as *mut T, step_size);
+        if mem::size_of::<T>() == 0 {
+            // SAFETY: due to unchecked casts of unsigned amounts to signed offsets the wraparound
+            // effectively results in unsigned pointers representing positions 0..usize::MAX,
+            // which is valid for ZSTs.
+            self.ptr = unsafe { arith_offset(self.ptr as *const i8, step_size as isize) as *mut T }
+        } else {
+            // SAFETY: the min() above ensures that step_size is in bounds
+            self.ptr = unsafe { self.ptr.add(step_size) };
+        }
+        // SAFETY: the min() above ensures that step_size is in bounds
+        unsafe {
+            ptr::drop_in_place(to_drop);
+        }
+        if step_size < n {
+            return Err(step_size);
+        }
+        Ok(())
+    }
+
+    #[inline]
     fn count(self) -> usize {
         self.len()
     }
@@ -203,6 +226,29 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
             Some(unsafe { ptr::read(self.end) })
         }
     }
+
+    #[inline]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        let step_size = self.len().min(n);
+        if mem::size_of::<T>() == 0 {
+            // SAFETY: same as for advance_by()
+            self.end = unsafe {
+                arith_offset(self.end as *const i8, step_size.wrapping_neg() as isize) as *mut T
+            }
+        } else {
+            // SAFETY: same as for advance_by()
+            self.end = unsafe { self.end.offset(step_size.wrapping_neg() as isize) };
+        }
+        let to_drop = ptr::slice_from_raw_parts_mut(self.end as *mut T, step_size);
+        // SAFETY: same as for advance_by()
+        unsafe {
+            ptr::drop_in_place(to_drop);
+        }
+        if step_size < n {
+            return Err(step_size);
+        }
+        Ok(())
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs
index cae4dae708e..c6159539b48 100644
--- a/library/alloc/tests/lib.rs
+++ b/library/alloc/tests/lib.rs
@@ -18,6 +18,7 @@
 #![feature(binary_heap_retain)]
 #![feature(binary_heap_as_slice)]
 #![feature(inplace_iteration)]
+#![feature(iter_advance_by)]
 #![feature(slice_group_by)]
 #![feature(slice_partition_dedup)]
 #![feature(vec_spare_capacity)]
diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs
index c2df50b48f5..00a878c0794 100644
--- a/library/alloc/tests/vec.rs
+++ b/library/alloc/tests/vec.rs
@@ -971,6 +971,24 @@ fn test_into_iter_leak() {
 }
 
 #[test]
+fn test_into_iter_advance_by() {
+    let mut i = vec![1, 2, 3, 4, 5].into_iter();
+    i.advance_by(0).unwrap();
+    i.advance_back_by(0).unwrap();
+    assert_eq!(i.as_slice(), [1, 2, 3, 4, 5]);
+
+    i.advance_by(1).unwrap();
+    i.advance_back_by(1).unwrap();
+    assert_eq!(i.as_slice(), [2, 3, 4]);
+
+    assert_eq!(i.advance_back_by(usize::MAX), Err(3));
+
+    assert_eq!(i.advance_by(usize::MAX), Err(0));
+
+    assert_eq!(i.len(), 0);
+}
+
+#[test]
 fn test_from_iter_specialization() {
     let src: Vec<usize> = vec![0usize; 1];
     let srcptr = src.as_ptr();