Rollup merge of #150230 - bend-n:trusted_len_spec_for_iter_next_chunk, r=Mark-Simulacrum

spec next chunk for trustedlen

relevant to rust-lang/rust#98326

library/core/src/array/mod.rs

@@ -11,7 +11,7 @@
 use crate::error::Error;
 use crate::hash::{self, Hash};
 use crate::intrinsics::transmute_unchecked;
-use crate::iter::{UncheckedIterator, repeat_n};
+use crate::iter::{TrustedLen, UncheckedIterator, repeat_n};
 use crate::marker::Destruct;
 use crate::mem::{self, ManuallyDrop, MaybeUninit};
 use crate::ops::{
@@ -1010,19 +1010,66 @@ fn drop(&mut self) {
 pub(crate) const fn iter_next_chunk<T, const N: usize>(
     iter: &mut impl [const] Iterator<Item = T>,
 ) -> Result<[T; N], IntoIter<T, N>> {
-    let mut array = [const { MaybeUninit::uninit() }; N];
-    let r = iter_next_chunk_erased(&mut array, iter);
-    match r {
-        Ok(()) => {
-            // SAFETY: All elements of `array` were populated.
-            Ok(unsafe { MaybeUninit::array_assume_init(array) })
-        }
-        Err(initialized) => {
-            // SAFETY: Only the first `initialized` elements were populated
-            Err(unsafe { IntoIter::new_unchecked(array, 0..initialized) })
+    iter.spec_next_chunk()
+}
+
+pub(crate) const trait SpecNextChunk<T, const N: usize>: Iterator<Item = T> {
+    fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter<T, N>>;
+}
+#[rustc_const_unstable(feature = "const_iter", issue = "92476")]
+impl<I: [const] Iterator<Item = T>, T, const N: usize> const SpecNextChunk<T, N> for I {
+    #[inline]
+    default fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter<T, N>> {
+        let mut array = [const { MaybeUninit::uninit() }; N];
+        let r = iter_next_chunk_erased(&mut array, self);
+        match r {
+            Ok(()) => {
+                // SAFETY: All elements of `array` were populated.
+                Ok(unsafe { MaybeUninit::array_assume_init(array) })
+            }
+            Err(initialized) => {
+                // SAFETY: Only the first `initialized` elements were populated
+                Err(unsafe { IntoIter::new_unchecked(array, 0..initialized) })
+            }
         }
     }
 }
+#[rustc_const_unstable(feature = "const_iter", issue = "92476")]
+impl<I: [const] Iterator<Item = T> + TrustedLen, T, const N: usize> const SpecNextChunk<T, N>
+    for I
+{
+    fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter<T, N>> {
+        let len = (*self).size_hint().0;
+        let mut array = [const { MaybeUninit::uninit() }; N];
+        if len < N {
+            // SAFETY: `TrustedLen`, an unsafe trait, requires that i can get len items out of it.
+            unsafe { write(&mut array, self, len) };
+            // SAFETY: Only the first `len` elements were populated
+            Err(unsafe { IntoIter::new_unchecked(array, 0..len) })
+        } else {
+            // SAFETY: `TrustedLen`, an unsafe trait, requires that i can get N items out of it.
+            unsafe { write(&mut array, self, N) };
+            // SAFETY: All N items were populated
+            Ok(unsafe { MaybeUninit::array_assume_init(array) })
+        }
+    }
+}
+// SAFETY: `from` must have len items, and len items must be < N.
+#[rustc_const_unstable(feature = "const_iter", issue = "92476")]
+const unsafe fn write<T, const N: usize>(
+    to: &mut [MaybeUninit<T>; N],
+    from: &mut impl [const] Iterator<Item = T>,
+    len: usize,
+) {
+    let mut guard = Guard { array_mut: to, initialized: 0 };
+    while guard.initialized < len {
+        // SAFETY: caller has guaranteed, from has len items.
+        let item = unsafe { from.next().unwrap_unchecked() };
+        // SAFETY: guard.initialized < len < N
+        unsafe { guard.push_unchecked(item) };
+    }
+    crate::mem::forget(guard);
+}
 
 /// Version of [`iter_next_chunk`] using a passed-in slice in order to avoid
 /// needing to monomorphize for every array length.

library/coretests/tests/iter/traits/iterator.rs

@@ -619,6 +619,23 @@ fn test_next_chunk() {
     assert_eq!(it.next_chunk::<0>().unwrap(), []);
 }
 
+#[test]
+fn test_next_chunk_untrusted() {
+    struct Untrusted<I: Iterator>(I);
+    impl<I: Iterator> Iterator for Untrusted<I> {
+        type Item = I::Item;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.0.next()
+        }
+    }
+    let mut it = Untrusted(0..12);
+    assert_eq!(it.next_chunk().unwrap(), [0, 1, 2, 3]);
+    assert_eq!(it.next_chunk().unwrap(), []);
+    assert_eq!(it.next_chunk().unwrap(), [4, 5, 6, 7, 8, 9]);
+    assert_eq!(it.next_chunk::<4>().unwrap_err().as_slice(), &[10, 11]);
+}
+
 #[test]
 fn test_collect_into_tuples() {
     let a = vec![(1, 2, 3), (4, 5, 6), (7, 8, 9)];