Rollup merge of #76635 - scottmcm:slice-as-chunks, r=LukasKalbertodt

Add [T]::as_chunks(_mut)

Allows getting the slices directly, rather than just through an iterator as in `array_chunks(_mut)`.  The constructors for those iterators are then written in terms of these methods, so the iterator constructors no longer have any `unsafe` of their own.

Unstable, of course. #74985

2 files changed, 71 insertions, 15 deletions

diff --git a/library/core/src/slice/iter.rs b/library/core/src/slice/iter.rs
index 24f955a70b6..e373936a6c7 100644
--- a/library/core/src/slice/iter.rs
+++ b/library/core/src/slice/iter.rs
@@ -2103,13 +2103,8 @@ pub struct ArrayChunks<'a, T: 'a, const N: usize> {
 impl<'a, T, const N: usize> ArrayChunks<'a, T, N> {
     #[inline]
     pub(super) fn new(slice: &'a [T]) -> Self {
-        let len = slice.len() / N;
-        let (fst, snd) = slice.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) };
-
-        Self { iter: array_slice.iter(), rem: snd }
+        let (array_slice, rem) = slice.as_chunks();
+        Self { iter: array_slice.iter(), rem }
     }
 
     /// Returns the remainder of the original slice that is not going to be
@@ -2230,14 +2225,8 @@ pub struct ArrayChunksMut<'a, T: 'a, const N: usize> {
 impl<'a, T, const N: usize> ArrayChunksMut<'a, T, N> {
     #[inline]
     pub(super) fn new(slice: &'a mut [T]) -> Self {
-        let len = slice.len() / N;
-        let (fst, snd) = slice.split_at_mut(len * N);
-        // SAFETY: We cast a slice of `len * N` elements into
-        // a slice of `len` many `N` elements chunks.
-        unsafe {
-            let array_slice: &mut [[T; N]] = from_raw_parts_mut(fst.as_mut_ptr().cast(), len);
-            Self { iter: array_slice.iter_mut(), rem: snd }
-        }
+        let (array_slice, rem) = slice.as_chunks_mut();
+        Self { iter: array_slice.iter_mut(), rem }
     }
 
     /// Returns the remainder of the original slice that is not going to be
diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs
index b6fd0c4986b..f514a90c2b8 100644
--- a/library/core/src/slice/mod.rs
+++ b/library/core/src/slice/mod.rs
@@ -882,6 +882,36 @@ impl<T> [T] {
         ChunksExactMut::new(self, chunk_size)
     }
 
+    /// Splits the slice into a slice of `N`-element arrays,
+    /// starting at the beginning of the slice,
+    /// and a remainder slice with length strictly less than `N`.
+    ///
+    /// # 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(slice_as_chunks)]
+    /// let slice = ['l', 'o', 'r', 'e', 'm'];
+    /// let (chunks, remainder) = slice.as_chunks();
+    /// assert_eq!(chunks, &[['l', 'o'], ['r', 'e']]);
+    /// assert_eq!(remainder, &['m']);
+    /// ```
+    #[unstable(feature = "slice_as_chunks", issue = "74985")]
+    #[inline]
+    pub fn as_chunks<const N: usize>(&self) -> (&[[T; N]], &[T]) {
+        assert_ne!(N, 0);
+        let len = self.len() / N;
+        let (multiple_of_n, remainder) = 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(multiple_of_n.as_ptr().cast(), len) };
+        (array_slice, remainder)
+    }
+
     /// Returns an iterator over `N` elements of the slice at a time, starting at the
     /// beginning of the slice.
     ///
@@ -916,6 +946,43 @@ impl<T> [T] {
         ArrayChunks::new(self)
     }
 
+    /// Splits the slice into a slice of `N`-element arrays,
+    /// starting at the beginning of the slice,
+    /// and a remainder slice with length strictly less than `N`.
+    ///
+    /// # 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(slice_as_chunks)]
+    /// let v = &mut [0, 0, 0, 0, 0];
+    /// let mut count = 1;
+    ///
+    /// let (chunks, remainder) = v.as_chunks_mut();
+    /// remainder[0] = 9;
+    /// for chunk in chunks {
+    ///     *chunk = [count; 2];
+    ///     count += 1;
+    /// }
+    /// assert_eq!(v, &[1, 1, 2, 2, 9]);
+    /// ```
+    #[unstable(feature = "slice_as_chunks", issue = "74985")]
+    #[inline]
+    pub fn as_chunks_mut<const N: usize>(&mut self) -> (&mut [[T; N]], &mut [T]) {
+        assert_ne!(N, 0);
+        let len = self.len() / N;
+        let (multiple_of_n, remainder) = self.split_at_mut(len * N);
+        let array_slice: &mut [[T; N]] =
+            // SAFETY: We cast a slice of `len * N` elements into
+            // a slice of `len` many `N` elements chunks.
+            unsafe { from_raw_parts_mut(multiple_of_n.as_mut_ptr().cast(), len) };
+        (array_slice, remainder)
+    }
+
     /// Returns an iterator over `N` elements of the slice at a time, starting at the
     /// beginning of the slice.
     ///