From 652f34d2709ec0c323ae632ba14992f80ebf7629 Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Sat, 12 Sep 2020 01:15:28 -0700
Subject: [PATCH] Add [T]::as_chunks_mut (as unstable)

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.
---
 library/core/src/slice/iter.rs | 19 ++--------
 library/core/src/slice/mod.rs  | 67 ++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 15 deletions(-)

diff --git a/library/core/src/slice/iter.rs b/library/core/src/slice/iter.rs
index 793cbf994956f..9c58153316aa5 100644
--- a/library/core/src/slice/iter.rs
+++ b/library/core/src/slice/iter.rs
@@ -2102,13 +2102,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
@@ -2229,14 +2224,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 b1ca093add581..d3591e8547525 100644
--- a/library/core/src/slice/mod.rs
+++ b/library/core/src/slice/mod.rs
@@ -883,6 +883,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.
     ///
@@ -917,6 +947,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.
     ///