diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index 065d071a2e3..69ca33cfc22 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -131,6 +131,7 @@ #![feature(trusted_len)] #![feature(trusted_random_access)] #![feature(try_trait_v2)] +#![feature(unchecked_math)] #![feature(unicode_internals)] #![feature(unsize)] // diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 0b62622de81..74bcac2b541 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -2274,6 +2274,51 @@ pub fn extend_from_within(&mut self, src: R) } } +impl Vec<[T; N], A> { + /// Takes a `Vec<[T; N]>` and flattens it into a `Vec`. + /// + /// # Panics + /// + /// Panics if the length of the resulting vector would overflow a `usize`. + /// + /// This is only possible when flattening a vector of arrays of zero-sized + /// types, and thus tends to be irrelevant in practice. If + /// `size_of::() > 0`, this will never panic. + /// + /// # Examples + /// + /// ``` + /// #![feature(slice_flatten)] + /// + /// let mut vec = vec![[1, 2, 3], [4, 5, 6], [7, 8, 9]]; + /// assert_eq!(vec.pop(), Some([7, 8, 9])); + /// + /// let mut flattened = vec.into_flattened(); + /// assert_eq!(flattened.pop(), Some(6)); + /// ``` + #[unstable(feature = "slice_flatten", issue = "95629")] + pub fn into_flattened(self) -> Vec { + let (ptr, len, cap, alloc) = self.into_raw_parts_with_alloc(); + let (new_len, new_cap) = if mem::size_of::() == 0 { + (len.checked_mul(N).expect("vec len overflow"), usize::MAX) + } else { + // SAFETY: + // - `cap * N` cannot overflow because the allocation is already in + // the address space. + // - Each `[T; N]` has `N` valid elements, so there are `len * N` + // valid elements in the allocation. + unsafe { (len.unchecked_mul(N), cap.unchecked_mul(N)) } + }; + // SAFETY: + // - `ptr` was allocated by `self` + // - `ptr` is well-aligned because `[T; N]` has the same alignment as `T`. + // - `new_cap` refers to the same sized allocation as `cap` because + // `new_cap * size_of::()` == `cap * size_of::<[T; N]>()` + // - `len` <= `cap`, so `len * N` <= `cap * N`. + unsafe { Vec::::from_raw_parts_in(ptr.cast(), new_len, new_cap, alloc) } + } +} + // This code generalizes `extend_with_{element,default}`. trait ExtendWith { fn next(&mut self) -> T; diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index abce47e5afe..16ceb8e373d 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -38,6 +38,7 @@ #![feature(const_str_from_utf8)] #![feature(nonnull_slice_from_raw_parts)] #![feature(panic_update_hook)] +#![feature(slice_flatten)] use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index 19e39ebf910..bc1397146dd 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -2408,3 +2408,10 @@ fn drop(&mut self) { assert_eq!(count.load(Ordering::SeqCst), 4); } + +#[test] +#[should_panic = "vec len overflow"] +fn test_into_flattened_size_overflow() { + let v = vec![[(); usize::MAX]; 2]; + let _ = v.into_flattened(); +} diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs index 77bf5f9dc34..e76c31635ea 100644 --- a/library/core/src/slice/mod.rs +++ b/library/core/src/slice/mod.rs @@ -3994,6 +3994,88 @@ pub fn take_last_mut<'a>(self: &mut &'a mut Self) -> Option<&'a mut T> { } } +#[cfg(not(bootstrap))] +impl [[T; N]] { + /// Takes a `&[[T; N]]`, and flattens it to a `&[T]`. + /// + /// # Panics + /// + /// This panics if the length of the resulting slice would overflow a `usize`. + /// + /// This is only possible when flattening a slice of arrays of zero-sized + /// types, and thus tends to be irrelevant in practice. If + /// `size_of::() > 0`, this will never panic. + /// + /// # Examples + /// + /// ``` + /// #![feature(slice_flatten)] + /// + /// assert_eq!([[1, 2, 3], [4, 5, 6]].flatten(), &[1, 2, 3, 4, 5, 6]); + /// + /// assert_eq!( + /// [[1, 2, 3], [4, 5, 6]].flatten(), + /// [[1, 2], [3, 4], [5, 6]].flatten(), + /// ); + /// + /// let slice_of_empty_arrays: &[[i32; 0]] = &[[], [], [], [], []]; + /// assert!(slice_of_empty_arrays.flatten().is_empty()); + /// + /// let empty_slice_of_arrays: &[[u32; 10]] = &[]; + /// assert!(empty_slice_of_arrays.flatten().is_empty()); + /// ``` + #[unstable(feature = "slice_flatten", issue = "95629")] + pub fn flatten(&self) -> &[T] { + let len = if crate::mem::size_of::() == 0 { + self.len().checked_mul(N).expect("slice len overflow") + } else { + // SAFETY: `self.len() * N` cannot overflow because `self` is + // already in the address space. + unsafe { self.len().unchecked_mul(N) } + }; + // SAFETY: `[T]` is layout-identical to `[T; N]` + unsafe { from_raw_parts(self.as_ptr().cast(), len) } + } + + /// Takes a `&mut [[T; N]]`, and flattens it to a `&mut [T]`. + /// + /// # Panics + /// + /// This panics if the length of the resulting slice would overflow a `usize`. + /// + /// This is only possible when flattening a slice of arrays of zero-sized + /// types, and thus tends to be irrelevant in practice. If + /// `size_of::() > 0`, this will never panic. + /// + /// # Examples + /// + /// ``` + /// #![feature(slice_flatten)] + /// + /// fn add_5_to_all(slice: &mut [i32]) { + /// for i in slice { + /// *i += 5; + /// } + /// } + /// + /// let mut array = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]; + /// add_5_to_all(array.flatten_mut()); + /// assert_eq!(array, [[6, 7, 8], [9, 10, 11], [12, 13, 14]]); + /// ``` + #[unstable(feature = "slice_flatten", issue = "95629")] + pub fn flatten_mut(&mut self) -> &mut [T] { + let len = if crate::mem::size_of::() == 0 { + self.len().checked_mul(N).expect("slice len overflow") + } else { + // SAFETY: `self.len() * N` cannot overflow because `self` is + // already in the address space. + unsafe { self.len().unchecked_mul(N) } + }; + // SAFETY: `[T]` is layout-identical to `[T; N]` + unsafe { from_raw_parts_mut(self.as_mut_ptr().cast(), len) } + } +} + trait CloneFromSpec { fn spec_clone_from(&mut self, src: &[T]); } diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 5f90a76ab74..e9955655295 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -93,6 +93,7 @@ #![feature(const_array_from_ref)] #![feature(const_slice_from_ref)] #![feature(waker_getters)] +#![feature(slice_flatten)] #![deny(unsafe_op_in_unsafe_fn)] extern crate test; diff --git a/library/core/tests/slice.rs b/library/core/tests/slice.rs index 798fcc3dfde..ada479147db 100644 --- a/library/core/tests/slice.rs +++ b/library/core/tests/slice.rs @@ -2504,3 +2504,19 @@ fn test_slice_from_ptr_range() { assert_eq!(slice::from_ptr_range(range), &arr); } } + +#[test] +#[cfg(not(bootstrap))] +#[should_panic = "slice len overflow"] +fn test_flatten_size_overflow() { + let x = &[[(); usize::MAX]; 2][..]; + let _ = x.flatten(); +} + +#[test] +#[cfg(not(bootstrap))] +#[should_panic = "slice len overflow"] +fn test_flatten_mut_size_overflow() { + let x = &mut [[(); usize::MAX]; 2][..]; + let _ = x.flatten_mut(); +}