diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs index dc34d2a21ee96..18ce212e38fd1 100644 --- a/library/core/src/slice/mod.rs +++ b/library/core/src/slice/mod.rs @@ -4664,6 +4664,124 @@ impl [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::() 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::()?` + let Some((head, rem)) = mem::replace(self, &mut []).split_first_chunk_mut::() 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::() 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::()?` + let Some((rem, tail)) = mem::replace(self, &mut []).split_last_chunk_mut::() 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 diff --git a/library/coretests/tests/lib.rs b/library/coretests/tests/lib.rs index 0387b442562db..8d6951568d22c 100644 --- a/library/coretests/tests/lib.rs +++ b/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)] diff --git a/library/coretests/tests/slice.rs b/library/coretests/tests/slice.rs index 6f60f71e8a477..b2de7f752c05b 100644 --- a/library/coretests/tests/slice.rs +++ b/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::() + ); + + 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::()); + } + { + // 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::() + ); + + 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::()); + } + { + // 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()];