Implement split_array and split_array_mut

5 files changed, 219 insertions, 0 deletions

diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs
index b27c36baf37..73340fda2cb 100644
--- a/library/core/src/array/mod.rs
+++ b/library/core/src/array/mod.rs
@@ -500,6 +500,84 @@ impl<T, const N: usize> [T; N] {
         // items.
         unsafe { collect_into_array_unchecked(&mut self.iter_mut()) }
     }
+
+    /// Divides one array reference into two at an index.
+    ///
+    /// The first will contain all indices from `[0, M)` (excluding
+    /// the index `M` itself) and the second will contain all
+    /// indices from `[M, N)` (excluding the index `N` itself).
+    ///
+    /// # Panics
+    ///
+    /// Panics if `M > N`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(split_array)]
+    ///
+    /// let v = [1, 2, 3, 4, 5, 6];
+    ///
+    /// {
+    ///    let (left, right) = v.split_array_ref::<0>();
+    ///    assert_eq!(left, &[]);
+    ///    assert_eq!(right, &[1, 2, 3, 4, 5, 6]);
+    /// }
+    ///
+    /// {
+    ///     let (left, right) = v.split_array_ref::<2>();
+    ///     assert_eq!(left, &[1, 2]);
+    ///     assert_eq!(right, &[3, 4, 5, 6]);
+    /// }
+    ///
+    /// {
+    ///     let (left, right) = v.split_array_ref::<6>();
+    ///     assert_eq!(left, &[1, 2, 3, 4, 5, 6]);
+    ///     assert_eq!(right, &[]);
+    /// }
+    /// ```
+    #[unstable(
+        feature = "split_array",
+        reason = "return type should have array as 2nd element",
+        issue = "90091"
+    )]
+    #[inline]
+    pub fn split_array_ref<const M: usize>(&self) -> (&[T; M], &[T]) {
+        (&self[..]).split_array_ref::<M>()
+    }
+
+    /// Divides one mutable array reference into two at an index.
+    ///
+    /// The first will contain all indices from `[0, M)` (excluding
+    /// the index `M` itself) and the second will contain all
+    /// indices from `[M, N)` (excluding the index `N` itself).
+    ///
+    /// # Panics
+    ///
+    /// Panics if `M > N`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(split_array)]
+    ///
+    /// let mut v = [1, 0, 3, 0, 5, 6];
+    /// let (left, right) = v.split_array_mut::<2>();
+    /// assert_eq!(left, &mut [1, 0][..]);
+    /// assert_eq!(right, &mut [3, 0, 5, 6]);
+    /// left[1] = 2;
+    /// right[1] = 4;
+    /// assert_eq!(v, [1, 2, 3, 4, 5, 6]);
+    /// ```
+    #[unstable(
+        feature = "split_array",
+        reason = "return type should have array as 2nd element",
+        issue = "90091"
+    )]
+    #[inline]
+    pub fn split_array_mut<const M: usize>(&mut self) -> (&mut [T; M], &mut [T]) {
+        (&mut self[..]).split_array_mut::<M>()
+    }
 }
 
 /// Pulls `N` items from `iter` and returns them as an array. If the iterator
diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs
index 664875a8773..aaa00653b99 100644
--- a/library/core/src/slice/mod.rs
+++ b/library/core/src/slice/mod.rs
@@ -1665,6 +1665,80 @@ impl<T> [T] {
         unsafe { (from_raw_parts_mut(ptr, mid), from_raw_parts_mut(ptr.add(mid), len - mid)) }
     }
 
+    /// Divides one slice into an array and a remainder slice at an index.
+    ///
+    /// The array will contain all indices from `[0, N)` (excluding
+    /// the index `N` itself) and the slice will contain all
+    /// indices from `[N, len)` (excluding the index `len` itself).
+    ///
+    /// # Panics
+    ///
+    /// Panics if `N > len`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(split_array)]
+    ///
+    /// let v = &[1, 2, 3, 4, 5, 6][..];
+    ///
+    /// {
+    ///    let (left, right) = v.split_array_ref::<0>();
+    ///    assert_eq!(left, &[]);
+    ///    assert_eq!(right, [1, 2, 3, 4, 5, 6]);
+    /// }
+    ///
+    /// {
+    ///     let (left, right) = v.split_array_ref::<2>();
+    ///     assert_eq!(left, &[1, 2]);
+    ///     assert_eq!(right, [3, 4, 5, 6]);
+    /// }
+    ///
+    /// {
+    ///     let (left, right) = v.split_array_ref::<6>();
+    ///     assert_eq!(left, &[1, 2, 3, 4, 5, 6]);
+    ///     assert_eq!(right, []);
+    /// }
+    /// ```
+    #[unstable(feature = "split_array", reason = "new API", issue = "90091")]
+    #[inline]
+    pub fn split_array_ref<const N: usize>(&self) -> (&[T; N], &[T]) {
+        let (a, b) = self.split_at(N);
+        // SAFETY: a points to [T; N]? Yes it's [T] of length N (checked by split_at)
+        unsafe { (&*(a.as_ptr() as *const [T; N]), b) }
+    }
+
+    /// Divides one mutable slice into an array and a remainder slice at an index.
+    ///
+    /// The array will contain all indices from `[0, N)` (excluding
+    /// the index `N` itself) and the slice will contain all
+    /// indices from `[N, len)` (excluding the index `len` itself).
+    ///
+    /// # Panics
+    ///
+    /// Panics if `N > len`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(split_array)]
+    ///
+    /// let mut v = &mut [1, 0, 3, 0, 5, 6][..];
+    /// let (left, right) = v.split_array_mut::<2>();
+    /// assert_eq!(left, &mut [1, 0]);
+    /// assert_eq!(right, [3, 0, 5, 6]);
+    /// left[1] = 2;
+    /// right[1] = 4;
+    /// assert_eq!(v, [1, 2, 3, 4, 5, 6]);
+    /// ```
+    #[unstable(feature = "split_array", reason = "new API", issue = "90091")]
+    #[inline]
+    pub fn split_array_mut<const N: usize>(&mut self) -> (&mut [T; N], &mut [T]) {
+        let (a, b) = self.split_at_mut(N);
+        // SAFETY: a points to [T; N]? Yes it's [T] of length N (checked by split_at_mut)
+        unsafe { (&mut *(a.as_mut_ptr() as *mut [T; N]), b) }
+    }
+
     /// Returns an iterator over subslices separated by elements that match
     /// `pred`. The matched element is not contained in the subslices.
     ///
diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs
index b3af1328c90..1d4307ca0fb 100644
--- a/library/core/tests/array.rs
+++ b/library/core/tests/array.rs
@@ -436,3 +436,36 @@ where
     std::panic::set_hook(prev_hook);
     result
 }
+
+#[test]
+fn array_split_array_mut() {
+    let mut v = [1, 2, 3, 4, 5, 6];
+
+    {
+        let (left, right) = v.split_array_mut::<0>();
+        assert_eq!(left, &mut []);
+        assert_eq!(right, &mut [1, 2, 3, 4, 5, 6]);
+    }
+
+    {
+        let (left, right) = v.split_array_mut::<6>();
+        assert_eq!(left, &mut [1, 2, 3, 4, 5, 6]);
+        assert_eq!(right, &mut []);
+    }
+}
+
+#[should_panic]
+#[test]
+fn array_split_array_ref_out_of_bounds() {
+    let v = [1, 2, 3, 4, 5, 6];
+
+    v.split_array_ref::<7>();
+}
+
+#[should_panic]
+#[test]
+fn array_split_array_mut_out_of_bounds() {
+    let mut v = [1, 2, 3, 4, 5, 6];
+
+    v.split_array_mut::<7>();
+}
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
index ab0295c6314..507a79c3de1 100644
--- a/library/core/tests/lib.rs
+++ b/library/core/tests/lib.rs
@@ -70,6 +70,7 @@
 #![feature(integer_atomics)]
 #![feature(int_roundings)]
 #![feature(slice_group_by)]
+#![feature(split_array)]
 #![feature(trusted_random_access)]
 #![feature(unsize)]
 #![feature(unzip_option)]
diff --git a/library/core/tests/slice.rs b/library/core/tests/slice.rs
index b6a326f3d73..7a008345243 100644
--- a/library/core/tests/slice.rs
+++ b/library/core/tests/slice.rs
@@ -2191,3 +2191,36 @@ mod swap_panics {
         x.swap(2, 5);
     }
 }
+
+#[test]
+fn slice_split_array_mut() {
+    let v = &mut [1, 2, 3, 4, 5, 6][..];
+
+    {
+        let (left, right) = v.split_array_mut::<0>();
+        assert_eq!(left, &mut []);
+        assert_eq!(right, [1, 2, 3, 4, 5, 6]);
+    }
+
+    {
+        let (left, right) = v.split_array_mut::<6>();
+        assert_eq!(left, &mut [1, 2, 3, 4, 5, 6]);
+        assert_eq!(right, []);
+    }
+}
+
+#[should_panic]
+#[test]
+fn slice_split_array_ref_out_of_bounds() {
+    let v = &[1, 2, 3, 4, 5, 6][..];
+
+    v.split_array_ref::<7>();
+}
+
+#[should_panic]
+#[test]
+fn slice_split_array_mut_out_of_bounds() {
+    let v = &mut [1, 2, 3, 4, 5, 6][..];
+
+    v.split_array_mut::<7>();
+}