Add methods to split off array chunks from slices

library/core/src/slice/mod.rs

@@ -4664,6 +4664,124 @@ impl<T> [T] {
         Some(last)
     }
 
+    /// Removes the first chunk of the slice and returns a reference
+    /// to it.
+    ///
+    /// Returns `None` if the slice has fewer than `N` elements.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_split_off_chunk)]
+    /// let mut slice: &[_] = &['a', 'b', 'c', 'd'];
+    /// let head = slice.split_off_first_chunk::<2>().unwrap();
+    ///
+    /// assert_eq!(slice, &['c', 'd']);
+    /// assert_eq!(head, &['a', 'b']);
+    /// ```
+    #[unstable(feature = "slice_split_off_chunk", issue = "none")]
+    #[rustc_const_unstable(feature = "const_slice_split_off_chunk", issue = "none")]
+    pub const fn split_off_first_chunk<'a, const N: usize>(
+        self: &mut &'a Self,
+    ) -> Option<&'a [T; N]> {
+        // FIXME(const-hack): Use `?` when available in const instead of `let-else`.
+        let Some((head, rem)) = self.split_first_chunk::<N>() else { return None };
+        *self = rem;
+        Some(head)
+    }
+
+    /// Removes the first chunk of the slice and returns a mutable
+    /// reference to it.
+    ///
+    /// Returns `None` if the slice has fewer than `N` elements.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_split_off_chunk)]
+    /// let mut slice: &mut [_] = &mut ['a', 'b', 'c', 'd'];
+    /// let head = slice.split_off_first_chunk_mut::<2>().unwrap();
+    ///
+    /// assert_eq!(slice, &mut ['c', 'd']);
+    /// assert_eq!(head, &mut ['a', 'b']);
+    /// ```
+    #[unstable(feature = "slice_split_off_chunk", issue = "none")]
+    #[rustc_const_unstable(feature = "const_slice_split_off_chunk", issue = "none")]
+    pub const fn split_off_first_chunk_mut<'a, const N: usize>(
+        self: &mut &'a mut Self,
+    ) -> Option<&'a mut [T; N]> {
+        if self.len() < N {
+            // Avoid failing in the next spot, which moves out of `self`.
+            return None;
+        }
+        // FIXME(const-hack): Use `mem::take` and `?` when available in const.
+        // Original: `mem::take(self).split_first_chunk_mut::<N>()?`
+        let Some((head, rem)) = mem::replace(self, &mut []).split_first_chunk_mut::<N>() else {
+            return None;
+        };
+        *self = rem;
+        Some(head)
+    }
+
+    /// Removes the last chunk of the slice and returns a reference
+    /// to it.
+    ///
+    /// Returns `None` if the slice has fewer than `N` elements.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_split_off_chunk)]
+    /// let mut slice: &[_] = &['a', 'b', 'c', 'd'];
+    /// let tail = slice.split_off_last_chunk::<2>().unwrap();
+    ///
+    /// assert_eq!(slice, &['a', 'b']);
+    /// assert_eq!(tail, &['c', 'd']);
+    /// ```
+    #[unstable(feature = "slice_split_off_chunk", issue = "none")]
+    #[rustc_const_unstable(feature = "const_slice_split_off_chunk", issue = "none")]
+    pub const fn split_off_last_chunk<'a, const N: usize>(
+        self: &mut &'a Self,
+    ) -> Option<&'a [T; N]> {
+        // FIXME(const-hack): Use `?` when available in const instead of `let-else`.
+        let Some((rem, tail)) = self.split_last_chunk::<N>() else { return None };
+        *self = rem;
+        Some(tail)
+    }
+
+    /// Removes the last chunk of the slice and returns a mutable
+    /// reference to it.
+    ///
+    /// Returns `None` if the slice has fewer than `N` elements.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(slice_split_off_chunk)]
+    /// let mut slice: &mut [_] = &mut ['a', 'b', 'c', 'd'];
+    /// let tail = slice.split_off_last_chunk_mut::<2>().unwrap();
+    ///
+    /// assert_eq!(slice, &mut ['a', 'b']);
+    /// assert_eq!(tail, &mut ['c', 'd']);
+    /// ```
+    #[unstable(feature = "slice_split_off_chunk", issue = "none")]
+    #[rustc_const_unstable(feature = "const_slice_split_off_chunk", issue = "none")]
+    pub const fn split_off_last_chunk_mut<'a, const N: usize>(
+        self: &mut &'a mut Self,
+    ) -> Option<&'a mut [T; N]> {
+        if self.len() < N {
+            // Avoid failing in the next spot, which moves out of `self`.
+            return None;
+        }
+        // FIXME(const-hack): Use `mem::take` and `?` when available in const.
+        // Original: `mem::take(self).split_last_chunk_mut::<N>()?`
+        let Some((rem, tail)) = mem::replace(self, &mut []).split_last_chunk_mut::<N>() else {
+            return None;
+        };
+        *self = rem;
+        Some(tail)
+    }
+
     /// Returns mutable references to many indices at once, without doing any checks.
     ///
     /// An index can be either a `usize`, a [`Range`] or a [`RangeInclusive`]. Note

library/coretests/tests/lib.rs

@@ -100,6 +100,7 @@
 #![feature(slice_index_methods)]
 #![feature(slice_internals)]
 #![feature(slice_partition_dedup)]
+#![feature(slice_split_off_chunk)]
 #![feature(slice_split_once)]
 #![feature(sliceindex_wrappers)]
 #![feature(split_array)]

library/coretests/tests/slice.rs

@@ -2333,6 +2333,87 @@ split_off_tests! {
     (split_off_mut_in_bounds_max_range_from, (usize::MAX..), Some(&mut [] as _), empty_max_mut!()),
 }
 
+#[test]
+fn test_split_off_chunk() {
+    let base_slice: &mut [u32] = &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+    let base_addr = core::ptr::from_ref(base_slice).addr();
+    {
+        let mut slice = &base_slice[..];
+        let empty_head = slice.split_off_first_chunk::<0>().unwrap();
+        assert_eq!(empty_head, &[]);
+        assert_eq!(slice, &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+        // Check that the address of the empty slice is correct.
+        assert_eq!(core::ptr::from_ref(empty_head).addr(), base_addr);
+
+        let head = slice.split_off_first_chunk::<2>().unwrap();
+        assert_eq!(head, &[1, 2]);
+        assert_eq!(slice, &[3, 4, 5, 6, 7, 8, 9, 10]);
+
+        let empty_tail = slice.split_off_last_chunk::<0>().unwrap();
+        assert_eq!(empty_tail, &[]);
+        assert_eq!(slice, &[3, 4, 5, 6, 7, 8, 9, 10]);
+        // Check that the address of the empty slice is correct.
+        assert_eq!(
+            core::ptr::from_ref(empty_tail).addr(),
+            base_addr + 10 * core::mem::size_of::<u32>()
+        );
+
+        let tail = slice.split_off_last_chunk::<2>().unwrap();
+        assert_eq!(tail, &[9, 10]);
+        assert_eq!(slice, &[3, 4, 5, 6, 7, 8]);
+
+        // Extract the whole slice and ensure the remainder is in the right spot.
+        let full_head = slice.split_off_first_chunk::<6>().unwrap();
+        assert_eq!(full_head, &[3, 4, 5, 6, 7, 8]);
+        assert_eq!(slice, &[]);
+        assert_eq!(core::ptr::from_ref(slice).addr(), base_addr + 8 * core::mem::size_of::<u32>());
+    }
+    {
+        // The same checks as above, but with mutable slices.
+        let mut slice = &mut base_slice[..];
+        let empty_head = slice.split_off_first_chunk_mut::<0>().unwrap();
+        assert_eq!(empty_head, &[]);
+        assert_eq!(slice, &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+        // Check that the address of the empty slice is correct.
+        assert_eq!(core::ptr::from_ref(empty_head).addr(), base_addr);
+
+        let head = slice.split_off_first_chunk_mut::<2>().unwrap();
+        assert_eq!(head, &mut [1, 2]);
+        assert_eq!(slice, &mut [3, 4, 5, 6, 7, 8, 9, 10]);
+
+        let empty_tail = slice.split_off_last_chunk_mut::<0>().unwrap();
+        assert_eq!(empty_tail, &mut []);
+        assert_eq!(slice, &mut [3, 4, 5, 6, 7, 8, 9, 10]);
+        // Check that the address of the empty slice is correct.
+        assert_eq!(
+            core::ptr::from_ref(empty_tail).addr(),
+            base_addr + 10 * core::mem::size_of::<u32>()
+        );
+
+        let tail = slice.split_off_last_chunk_mut::<2>().unwrap();
+        assert_eq!(tail, &mut [9, 10]);
+        assert_eq!(slice, &mut [3, 4, 5, 6, 7, 8]);
+
+        // Extract the whole slice and ensure the remainder is in the right spot.
+        let full_head = slice.split_off_first_chunk_mut::<6>().unwrap();
+        assert_eq!(full_head, &mut [3, 4, 5, 6, 7, 8]);
+        assert_eq!(slice, &mut []);
+        assert_eq!(core::ptr::from_ref(slice).addr(), base_addr + 8 * core::mem::size_of::<u32>());
+    }
+    {
+        // Check that oversized chunks return `None` and don't modify the original slice.
+        let mut slice = &base_slice[..];
+        assert!(slice.split_off_first_chunk::<12>().is_none());
+        assert!(slice.split_off_last_chunk::<12>().is_none());
+        assert_eq!(slice, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+
+        let mut slice = &mut base_slice[..];
+        assert!(slice.split_off_first_chunk_mut::<12>().is_none());
+        assert!(slice.split_off_last_chunk_mut::<12>().is_none());
+        assert_eq!(slice, &mut [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+    }
+}
+
 #[test]
 fn test_slice_from_ptr_range() {
     let arr = ["foo".to_owned(), "bar".to_owned()];