Auto merge of #74373 - lcnr:array_chunks, r=withoutboats

add `slice::array_chunks` to std

Now that #74113 has landed, these methods are suddenly usable. A rebirth of #72334

Tests are directly copied from `chunks_exact` and some additional tests for type inference.

r? @withoutboats as you are both part of t-libs and working on const generics. closes #60735

library/alloc/src/lib.rs

@@ -75,6 +75,7 @@
 #![cfg_attr(not(test), feature(generator_trait))]
 #![cfg_attr(test, feature(test))]
 #![feature(allocator_api)]
+#![feature(array_chunks)]
 #![feature(allow_internal_unstable)]
 #![feature(arbitrary_self_types)]
 #![feature(box_patterns)]

library/alloc/src/slice.rs

@@ -95,6 +95,8 @@ use crate::borrow::ToOwned;
 use crate::boxed::Box;
 use crate::vec::Vec;
 
+#[unstable(feature = "array_chunks", issue = "74985")]
+pub use core::slice::ArrayChunks;
 #[stable(feature = "slice_get_slice", since = "1.28.0")]
 pub use core::slice::SliceIndex;
 #[stable(feature = "from_ref", since = "1.28.0")]

library/core/src/slice/mod.rs

@@ -680,7 +680,7 @@ impl<T> [T] {
     #[stable(feature = "rust1", since = "1.0.0")]
     #[inline]
     pub fn windows(&self, size: usize) -> Windows<'_, T> {
-        assert!(size != 0);
+        assert_ne!(size, 0);
         Windows { v: self, size }
     }
 
@@ -714,7 +714,7 @@ impl<T> [T] {
     #[stable(feature = "rust1", since = "1.0.0")]
     #[inline]
     pub fn chunks(&self, chunk_size: usize) -> Chunks<'_, T> {
-        assert!(chunk_size != 0);
+        assert_ne!(chunk_size, 0);
         Chunks { v: self, chunk_size }
     }
 
@@ -752,7 +752,7 @@ impl<T> [T] {
     #[stable(feature = "rust1", since = "1.0.0")]
     #[inline]
     pub fn chunks_mut(&mut self, chunk_size: usize) -> ChunksMut<'_, T> {
-        assert!(chunk_size != 0);
+        assert_ne!(chunk_size, 0);
         ChunksMut { v: self, chunk_size }
     }
 
@@ -789,7 +789,7 @@ impl<T> [T] {
     #[stable(feature = "chunks_exact", since = "1.31.0")]
     #[inline]
     pub fn chunks_exact(&self, chunk_size: usize) -> ChunksExact<'_, T> {
-        assert!(chunk_size != 0);
+        assert_ne!(chunk_size, 0);
         let rem = self.len() % chunk_size;
         let len = self.len() - rem;
         let (fst, snd) = self.split_at(len);
@@ -834,13 +834,52 @@ impl<T> [T] {
     #[stable(feature = "chunks_exact", since = "1.31.0")]
     #[inline]
     pub fn chunks_exact_mut(&mut self, chunk_size: usize) -> ChunksExactMut<'_, T> {
-        assert!(chunk_size != 0);
+        assert_ne!(chunk_size, 0);
         let rem = self.len() % chunk_size;
         let len = self.len() - rem;
         let (fst, snd) = self.split_at_mut(len);
         ChunksExactMut { v: fst, rem: snd, chunk_size }
     }
 
+    /// Returns an iterator over `N` elements of the slice at a time, starting at the
+    /// beginning of the slice.
+    ///
+    /// The chunks are slices and do not overlap. If `N` does not divide the length of the
+    /// slice, then the last up to `N-1` elements will be omitted and can be retrieved
+    /// from the `remainder` function of the iterator.
+    ///
+    /// This method is the const generic equivalent of [`chunks_exact`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if `N` is 0. This check will most probably get changed to a compile time
+    /// error before this method gets stabilized.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(array_chunks)]
+    /// let slice = ['l', 'o', 'r', 'e', 'm'];
+    /// let mut iter = slice.array_chunks();
+    /// assert_eq!(iter.next().unwrap(), &['l', 'o']);
+    /// assert_eq!(iter.next().unwrap(), &['r', 'e']);
+    /// assert!(iter.next().is_none());
+    /// assert_eq!(iter.remainder(), &['m']);
+    /// ```
+    ///
+    /// [`chunks_exact`]: #method.chunks_exact
+    #[unstable(feature = "array_chunks", issue = "74985")]
+    #[inline]
+    pub fn array_chunks<const N: usize>(&self) -> ArrayChunks<'_, T, N> {
+        assert_ne!(N, 0);
+        let len = self.len() / N;
+        let (fst, snd) = self.split_at(len * N);
+        // SAFETY: We cast a slice of `len * N` elements into
+        // a slice of `len` many `N` elements chunks.
+        let array_slice: &[[T; N]] = unsafe { from_raw_parts(fst.as_ptr().cast(), len) };
+        ArrayChunks { iter: array_slice.iter(), rem: snd }
+    }
+
     /// Returns an iterator over `chunk_size` elements of the slice at a time, starting at the end
     /// of the slice.
     ///
@@ -5432,6 +5471,110 @@ unsafe impl<'a, T> TrustedRandomAccess for ChunksExactMut<'a, T> {
     }
 }
 
+/// An iterator over a slice in (non-overlapping) chunks (`N` elements at a
+/// time), starting at the beginning of the slice.
+///
+/// When the slice len is not evenly divided by the chunk size, the last
+/// up to `chunk_size-1` elements will be omitted but can be retrieved from
+/// the [`remainder`] function from the iterator.
+///
+/// This struct is created by the [`array_chunks`] method on [slices].
+///
+/// [`array_chunks`]: ../../std/primitive.slice.html#method.array_chunks
+/// [`remainder`]: ../../std/slice/struct.ArrayChunks.html#method.remainder
+/// [slices]: ../../std/primitive.slice.html
+#[derive(Debug)]
+#[unstable(feature = "array_chunks", issue = "74985")]
+pub struct ArrayChunks<'a, T: 'a, const N: usize> {
+    iter: Iter<'a, [T; N]>,
+    rem: &'a [T],
+}
+
+impl<'a, T, const N: usize> ArrayChunks<'a, T, N> {
+    /// Returns the remainder of the original slice that is not going to be
+    /// returned by the iterator. The returned slice has at most `chunk_size-1`
+    /// elements.
+    #[unstable(feature = "array_chunks", issue = "74985")]
+    pub fn remainder(&self) -> &'a [T] {
+        self.rem
+    }
+}
+
+// FIXME(#26925) Remove in favor of `#[derive(Clone)]`
+#[unstable(feature = "array_chunks", issue = "74985")]
+impl<T, const N: usize> Clone for ArrayChunks<'_, T, N> {
+    fn clone(&self) -> Self {
+        ArrayChunks { iter: self.iter.clone(), rem: self.rem }
+    }
+}
+
+#[unstable(feature = "array_chunks", issue = "74985")]
+impl<'a, T, const N: usize> Iterator for ArrayChunks<'a, T, N> {
+    type Item = &'a [T; N];
+
+    #[inline]
+    fn next(&mut self) -> Option<&'a [T; N]> {
+        self.iter.next()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.count()
+    }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        self.iter.nth(n)
+    }
+
+    #[inline]
+    fn last(self) -> Option<Self::Item> {
+        self.iter.last()
+    }
+}
+
+#[unstable(feature = "array_chunks", issue = "74985")]
+impl<'a, T, const N: usize> DoubleEndedIterator for ArrayChunks<'a, T, N> {
+    #[inline]
+    fn next_back(&mut self) -> Option<&'a [T; N]> {
+        self.iter.next_back()
+    }
+
+    #[inline]
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+        self.iter.nth_back(n)
+    }
+}
+
+#[unstable(feature = "array_chunks", issue = "74985")]
+impl<T, const N: usize> ExactSizeIterator for ArrayChunks<'_, T, N> {
+    fn is_empty(&self) -> bool {
+        self.iter.is_empty()
+    }
+}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<T, const N: usize> TrustedLen for ArrayChunks<'_, T, N> {}
+
+#[unstable(feature = "array_chunks", issue = "74985")]
+impl<T, const N: usize> FusedIterator for ArrayChunks<'_, T, N> {}
+
+#[doc(hidden)]
+#[unstable(feature = "array_chunks", issue = "74985")]
+unsafe impl<'a, T, const N: usize> TrustedRandomAccess for ArrayChunks<'a, T, N> {
+    unsafe fn get_unchecked(&mut self, i: usize) -> &'a [T; N] {
+        unsafe { self.iter.get_unchecked(i) }
+    }
+    fn may_have_side_effect() -> bool {
+        false
+    }
+}
+
 /// An iterator over a slice in (non-overlapping) chunks (`chunk_size` elements at a
 /// time), starting at the end of the slice.
 ///

library/core/tests/lib.rs

@@ -1,4 +1,5 @@
 #![feature(alloc_layout_extra)]
+#![feature(array_chunks)]
 #![feature(bool_to_option)]
 #![feature(bound_cloned)]
 #![feature(box_syntax)]

library/core/tests/slice.rs

@@ -473,6 +473,97 @@ fn test_chunks_exact_mut_zip() {
     assert_eq!(v1, [13, 14, 19, 20, 4]);
 }
 
+#[test]
+fn test_array_chunks_infer() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, -4];
+    let c = v.array_chunks();
+    for &[a, b, c] in c {
+        assert_eq!(a + b + c, 3);
+    }
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4, 5, 6];
+    let total = v2.array_chunks().map(|&[a, b]| a * b).sum::<i32>();
+    assert_eq!(total, 2 * 3 + 4 * 5);
+}
+
+#[test]
+fn test_array_chunks_count() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let c = v.array_chunks::<3>();
+    assert_eq!(c.count(), 2);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let c2 = v2.array_chunks::<2>();
+    assert_eq!(c2.count(), 2);
+
+    let v3: &[i32] = &[];
+    let c3 = v3.array_chunks::<2>();
+    assert_eq!(c3.count(), 0);
+}
+
+#[test]
+fn test_array_chunks_nth() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let mut c = v.array_chunks::<2>();
+    assert_eq!(c.nth(1).unwrap(), &[2, 3]);
+    assert_eq!(c.next().unwrap(), &[4, 5]);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4, 5, 6];
+    let mut c2 = v2.array_chunks::<3>();
+    assert_eq!(c2.nth(1).unwrap(), &[3, 4, 5]);
+    assert_eq!(c2.next(), None);
+}
+
+#[test]
+fn test_array_chunks_nth_back() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let mut c = v.array_chunks::<2>();
+    assert_eq!(c.nth_back(1).unwrap(), &[2, 3]);
+    assert_eq!(c.next().unwrap(), &[0, 1]);
+    assert_eq!(c.next(), None);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let mut c2 = v2.array_chunks::<3>();
+    assert_eq!(c2.nth_back(0).unwrap(), &[0, 1, 2]);
+    assert_eq!(c2.next(), None);
+    assert_eq!(c2.next_back(), None);
+
+    let v3: &[i32] = &[0, 1, 2, 3, 4];
+    let mut c3 = v3.array_chunks::<10>();
+    assert_eq!(c3.nth_back(0), None);
+}
+
+#[test]
+fn test_array_chunks_last() {
+    let v: &[i32] = &[0, 1, 2, 3, 4, 5];
+    let c = v.array_chunks::<2>();
+    assert_eq!(c.last().unwrap(), &[4, 5]);
+
+    let v2: &[i32] = &[0, 1, 2, 3, 4];
+    let c2 = v2.array_chunks::<2>();
+    assert_eq!(c2.last().unwrap(), &[2, 3]);
+}
+
+#[test]
+fn test_array_chunks_remainder() {
+    let v: &[i32] = &[0, 1, 2, 3, 4];
+    let c = v.array_chunks::<2>();
+    assert_eq!(c.remainder(), &[4]);
+}
+
+#[test]
+fn test_array_chunks_zip() {
+    let v1: &[i32] = &[0, 1, 2, 3, 4];
+    let v2: &[i32] = &[6, 7, 8, 9, 10];
+
+    let res = v1
+        .array_chunks::<2>()
+        .zip(v2.array_chunks::<2>())
+        .map(|(a, b)| a.iter().sum::<i32>() + b.iter().sum::<i32>())
+        .collect::<Vec<_>>();
+    assert_eq!(res, vec![14, 22]);
+}
+
 #[test]
 fn test_rchunks_count() {
     let v: &[i32] = &[0, 1, 2, 3, 4, 5];