diff --git a/library/core/src/array/guard.rs b/library/core/src/array/guard.rs new file mode 100644 index 0000000000000..3077a37e7a77f --- /dev/null +++ b/library/core/src/array/guard.rs @@ -0,0 +1,96 @@ +use crate::mem::{forget, replace, MaybeUninit}; +use crate::ptr; + +/// The internal-use drop guard for implementing array methods. +/// +/// This is free to be changed whenever. Its purpose is not to provide a +/// beautiful safe interface, but to make the unsafe details of `super`'s +/// other methods slightly more obvious and have reduced code duplication. +pub struct Guard<'a, T, const N: usize> { + array_mut: &'a mut [MaybeUninit; N], + initialized: usize, +} + +impl<'a, T, const N: usize> Guard<'a, T, N> { + #[inline] + pub fn new(buffer: &'a mut [MaybeUninit; N]) -> Self { + Self { array_mut: buffer, initialized: 0 } + } + + #[inline] + pub fn len(&self) -> usize { + self.initialized + } + + /// Initialize the next item + /// + /// # Safety + /// + /// Requires `self.len() < N`. + #[inline] + pub unsafe fn push_unchecked(&mut self, value: T) { + debug_assert!(self.len() < N); + // SAFETY: The precondition means we have space + unsafe { + self.array_mut.get_unchecked_mut(self.initialized).write(value); + } + self.initialized += 1; + } + + /// Initialize the next `CHUNK` item(s) + /// + /// # Safety + /// + /// Requires `self.len() + CHUNK <= N`. + #[inline] + pub unsafe fn push_chunk_unchecked(&mut self, chunk: [T; CHUNK]) { + assert!(CHUNK <= N); + debug_assert!(N - self.len() >= CHUNK); + // SAFETY: The precondition means we have space + unsafe { + // Since we're going to write multiple items, make sure not to do so + // via a `&mut MaybeUninit`, as that would violate stacked borrows. + let first = self.array_mut.as_mut_ptr(); + let p = first.add(self.initialized).cast(); + ptr::write(p, chunk); + } + self.initialized += CHUNK; + } + + /// Read the whole buffer as an initialized array. + /// + /// This always de-initializes the original buffer -- even if `T: Copy`. + /// + /// # Safety + /// + /// Requires `self.len() == N`. + #[inline] + pub unsafe fn into_array_unchecked(self) -> [T; N] { + debug_assert_eq!(self.len(), N); + + // This tells LLVM and MIRI that we don't care about the buffer after, + // and the extra `undef` write is trivial for it to optimize away. + let buffer = replace(self.array_mut, MaybeUninit::uninit_array()); + + // SAFETY: the condition above asserts that all elements are + // initialized. + let out = unsafe { MaybeUninit::array_assume_init(buffer) }; + + forget(self); + + out + } +} + +impl Drop for Guard<'_, T, N> { + fn drop(&mut self) { + debug_assert!(self.initialized <= N); + + // SAFETY: this slice will contain only initialized objects. + unsafe { + crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut( + &mut self.array_mut.get_unchecked_mut(..self.initialized), + )); + } + } +} diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs index fe7b3576e2f5f..688c759175cd0 100644 --- a/library/core/src/array/iter.rs +++ b/library/core/src/array/iter.rs @@ -104,6 +104,27 @@ impl IntoIter { MaybeUninit::slice_assume_init_mut(slice) } } + + /// Returns the remaining `M` items of this iterator as an array. + /// + /// If there are more than `M` items left in this iterator, the rest of + /// them will be leaked. So that should be avoided, but it's sound. + /// + /// # Safety + /// + /// There must be at least `M` items remaining. + pub(crate) unsafe fn into_array_unchecked(self) -> [T; M] { + debug_assert!(self.len() >= M); + + // SAFETY: The precondition of at least M items left means that + // there are enough valid contiguous items for the `read`. + let array: [T; M] = unsafe { ptr::read(self.as_slice().as_ptr().cast()) }; + + // We better not run any drops for the items we're about to return. + mem::forget(self); + + array + } } #[stable(feature = "array_value_iter_impls", since = "1.40.0")] diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 23fd1453e54d1..1aa22c48049a2 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -17,8 +17,11 @@ use crate::ops::{ use crate::slice::{Iter, IterMut}; mod equality; +mod guard; mod iter; +use guard::Guard; + #[stable(feature = "array_value_iter", since = "1.51.0")] pub use iter::IntoIter; @@ -94,6 +97,97 @@ where unsafe { try_collect_into_array_unchecked(&mut (0..N).map(cb)) } } +/// Creates an array by repeatedly cloning a value. +/// +/// This is only needed for types which are not `Copy`. If `T` is always `Copy`, +/// you can use `[value; N]` instead. +/// +/// If you don't have a constant `N`, but instead of runtime value `n`, then you +/// can create a `Vec` (instead of an array) via `vec![value; n]`. +/// +/// # Examples +/// +/// ``` +/// #![feature(array_repeat)] +/// +/// let array: [_; 4] = std::array::repeat(vec![1, 2, 3, 4]); +/// assert_eq!( +/// array, +/// [ +/// vec![1, 2, 3, 4], +/// vec![1, 2, 3, 4], +/// vec![1, 2, 3, 4], +/// vec![1, 2, 3, 4], +/// ] +/// ); +/// ``` +#[inline] +#[unstable(feature = "array_repeat", issue = "91613")] +pub fn repeat(value: T) -> [T; N] { + SpecArrayFill::repeat(value) +} + +/// Creates an array by repeatedly calling a closure. +/// +/// The elements will appear in the array in the order returned. +/// +/// # Examples +/// +/// ``` +/// #![feature(array_repeat)] +/// +/// let array: [Vec; 4] = std::array::repeat_with(Vec::new); +/// assert_eq!(array, [vec![], vec![], vec![], vec![]]); +/// +/// let mut p = 1; +/// let array: [_; 4] = std::array::repeat_with(|| { p *= 2; p }); +/// assert_eq!(array, [2, 4, 8, 16]); +/// ``` +#[inline] +#[unstable(feature = "array_repeat", issue = "91613")] +pub fn repeat_with(f: impl FnMut() -> T) -> [T; N] { + // SAFETY: `iter::repeat_with` is infinite, so is clearly long enough. + unsafe { collect_into_array_unchecked(&mut crate::iter::repeat_with(f)) } +} + +/// Creates an array by repeatedly calling a fallibly closure. +/// +/// If all invocations succeed, then this will return an array of those values. +/// Otherwise it will immediately return the first error encountered. +/// +/// The elements will appear in the array in the order returned. +/// +/// The return type of this function depends on the return type of the closure. +/// If the closure returns `Result`, then this function will return +/// `Result<[T; N], E>`. If the closure returns `Option`, then this function +/// will return `Option<[T; N]>`. +/// +/// # Examples +/// +/// ``` +/// #![feature(array_repeat)] +/// +/// let mut p = 1_u8; +/// let mut f = || -> Option { p = p.checked_mul(2)?; Some(p) }; +/// let array: Option<[_; 4]> = std::array::try_repeat_with(&mut f); +/// assert_eq!(array, Some([2, 4, 8, 16])); +/// let array: Option<[_; 2]> = std::array::try_repeat_with(&mut f); +/// assert_eq!(array, Some([32, 64])); +/// let array: Option<[_; 4]> = std::array::try_repeat_with(&mut f); +/// assert_eq!(array, None); +/// ``` +#[inline] +#[unstable(feature = "array_repeat", issue = "91613")] +pub fn try_repeat_with( + f: impl FnMut() -> R, +) -> ChangeOutputType +where + R::Residual: Residual<[R::Output; N]>, +{ + // SAFETY: `iter::repeat_with` is infinite, so is clearly long enough. + unsafe { try_collect_into_array_unchecked(&mut crate::iter::repeat_with(f)) } +} + /// Converts a reference to `T` into a reference to an array of length 1 (without copying). #[stable(feature = "array_from_ref", since = "1.53.0")] #[rustc_const_unstable(feature = "const_array_from_ref", issue = "90206")] @@ -473,6 +567,105 @@ impl [T; N] { unsafe { try_collect_into_array_unchecked(&mut IntoIterator::into_iter(self).map(f)) } } + /// Creates a new array with length equal to `NEW_LEN`. + /// + /// If `NEW_LEN` is greater than `N`, the additional elements are filled + /// with `value`. If `NEW_LEN` is less than `N`, the array is truncated. + /// + /// This method requires `T` to implement [`Clone`], in order to be able to + /// clone the passed value. If you need more flexibility (or want to rely + /// on [`Default`] instead of [`Clone`]), use [`resize_with`](#method.resize_with). + /// + /// # Examples + /// + /// ``` + /// #![feature(array_resize)] + /// + /// let array = ["hello"]; + /// let array = array.resize::<3>("world"); + /// assert_eq!(array, ["hello", "world", "world"]); + /// + /// let array = [1, 2, 3, 4]; + /// let array: [i32; 2] = array.resize(0); + /// assert_eq!(array, [1, 2]); + /// + /// let rgb = [0xA8, 0x3C, 0x09]; + /// let pixel = u32::from_be_bytes(rgb.resize(0)); + /// assert_eq!(pixel, 0xA83C0900); + /// ``` + #[must_use = "Unlike `Vec::resize`, this returns a new array"] + #[unstable(feature = "array_resize", issue = "91615")] + pub fn resize(self, value: T) -> [T; NEW_LEN] + where + T: Clone, + { + SpecArrayFill::resize(self, value) + } + + /// Creates a new array with length equal to `NEW_LEN`. + /// + /// If `NEW_LEN` is greater than `N`, the additional elements are filled + /// by calling `f`. The return values from `f` will end up in the array + /// in the order they have been generated. + /// + /// If `NEW_LEN` is less than `N`, the array is truncated. + /// + /// This method uses a closure to create new values. If you'd rather `Clone` + /// a given value, use [`resize`](#method.resize). If you want to use the + /// [`Default`] trait to generate values, you can pass [`Default::default`]. + /// + /// # Examples + /// + /// ``` + /// #![feature(array_resize)] + /// + /// let array = [1, 2, 3]; + /// let array: [_; 5] = array.resize_with(Default::default); + /// assert_eq!(array, [1, 2, 3, 0, 0]); + /// let array: [_; 4] = array.resize_with(Default::default); + /// assert_eq!(array, [1, 2, 3, 0]); + /// + /// let array = [7]; + /// let mut p = 1; + /// let array: [_; 5] = array.resize_with(|| { p *= 2; p }); + /// assert_eq!(array, [7, 2, 4, 8, 16]); + /// ``` + #[must_use = "Unlike `Vec::resize_with`, this returns a new array"] + #[unstable(feature = "array_resize", issue = "91615")] + pub fn resize_with(self, mut f: impl FnMut() -> T) -> [T; NEW_LEN] { + if NEW_LEN <= N { + self.truncate() + } else { + let mut buffer = MaybeUninit::uninit_array::(); + let mut guard = Guard::new(&mut buffer); + + // SAFETY: since `NEW_LEN > N` and we haven't pushed anything into + // the guard yet, there's definitely enough space. + unsafe { guard.push_chunk_unchecked(self) }; + + while guard.len() < NEW_LEN { + // SAFETY: The loop condition ensures there's space available + unsafe { guard.push_unchecked(f()) }; + } + + // SAFETY: The previous loop filled all the remaining spots. + unsafe { guard.into_array_unchecked() } + } + } + + // Leaving this an implementation detail for now, since it wants to enforce + // `M <= N` at compile-time, but that's unavailable on stable. + fn truncate(self) -> [T; M] { + assert!(M <= N); + + let remove = N - M; + let mut iter = IntoIterator::into_iter(self); + iter.advance_back_by(remove).expect("in-bounds by construction"); + + // SAFETY: `N - (N - M) => M`, so this is exactly the correct length + unsafe { iter.into_array_unchecked() } + } + /// 'Zips up' two arrays into a single array of pairs. /// /// `zip()` returns a new array where every element is a tuple where the @@ -647,6 +840,80 @@ impl [T; N] { } } +// Like `SpecFill` over in the `slice` module +trait SpecArrayFill: Clone { + fn repeat(value: Self) -> [Self; N]; + fn resize(array: [Self; N], value: Self) -> [Self; M]; + fn fill_rest(guard: Guard<'_, Self, N>, value: Self) -> [Self; N]; +} + +impl SpecArrayFill for T { + default fn repeat(value: Self) -> [Self; N] { + // Not using `from_fn` or `collect_into_array_unchecked` so that we can + // *move* the argument into the final element instead of always cloning. + + let mut buffer = MaybeUninit::uninit_array::(); + let guard = Guard::new(&mut buffer); + Self::fill_rest(guard, value) + } + + default fn resize(array: [Self; N], value: Self) -> [Self; M] { + if M <= N { + array.truncate() + } else { + let mut buffer = MaybeUninit::uninit_array::(); + let mut guard = Guard::new(&mut buffer); + + // SAFETY: since `M > N` and we haven't pushed anything into the + // guard yet, there's definitely enough space. + unsafe { guard.push_chunk_unchecked(array) }; + + Self::fill_rest(guard, value) + } + } + + default fn fill_rest(mut guard: Guard<'_, T, N>, value: Self) -> [Self; N] { + if guard.len() != N { + while N - guard.len() >= 2 { + // SAFETY: We checked there's at least two spaces + unsafe { guard.push_unchecked(value.clone()) }; + } + + // SAFETY: The `len != N` check ensures that we needed to write + // at least one element, and the loop only writes when two or more + // are needed, so there's always exactly one spot left here. + unsafe { guard.push_unchecked(value) }; + } + + // SAFETY: The previous pushes filled all the remaining spots. + unsafe { guard.into_array_unchecked() } + } +} + +// No need for special handling of the last item with `Copy` types. +// (Indeed, it can make it harder for the `memset` to kick in.) +// So specialize them to do the simpler thing instead. +impl SpecArrayFill for T { + fn repeat(value: Self) -> [Self; N] { + // Don't generate the closure type for every `N` + fn cloner(x: T) -> impl Fn(usize) -> T { + move |_| x + } + + from_fn(cloner(value)) + } + + fn fill_rest(mut guard: Guard<'_, T, N>, value: Self) -> [Self; N] { + while guard.len() < N { + // SAFETY: The loop condition ensures there's space available + unsafe { guard.push_unchecked(value) }; + } + + // SAFETY: The previous pushes filled all the remaining spots. + unsafe { guard.into_array_unchecked() } + } +} + /// Pulls `N` items from `iter` and returns them as an array. If the iterator /// yields fewer than `N` items, this function exhibits undefined behavior. /// @@ -710,26 +977,8 @@ where return unsafe { Some(Try::from_output(mem::zeroed())) }; } - struct Guard<'a, T, const N: usize> { - array_mut: &'a mut [MaybeUninit; N], - initialized: usize, - } - - impl Drop for Guard<'_, T, N> { - fn drop(&mut self) { - debug_assert!(self.initialized <= N); - - // SAFETY: this slice will contain only initialized objects. - unsafe { - crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut( - &mut self.array_mut.get_unchecked_mut(..self.initialized), - )); - } - } - } - let mut array = MaybeUninit::uninit_array::(); - let mut guard = Guard { array_mut: &mut array, initialized: 0 }; + let mut guard = Guard::new(&mut array); while let Some(item_rslt) = iter.next() { let item = match item_rslt.branch() { @@ -743,17 +992,14 @@ where // loop and the loop is aborted once it reaches N (which is // `array.len()`). unsafe { - guard.array_mut.get_unchecked_mut(guard.initialized).write(item); + guard.push_unchecked(item); } - guard.initialized += 1; // Check if the whole array was initialized. - if guard.initialized == N { - mem::forget(guard); - + if guard.len() == N { // SAFETY: the condition above asserts that all elements are // initialized. - let out = unsafe { MaybeUninit::array_assume_init(array) }; + let out = unsafe { guard.into_array_unchecked() }; return Some(Try::from_output(out)); } } diff --git a/src/test/codegen/array-methods.rs b/src/test/codegen/array-methods.rs new file mode 100644 index 0000000000000..a2bdcaf956f6c --- /dev/null +++ b/src/test/codegen/array-methods.rs @@ -0,0 +1,68 @@ +// no-system-llvm +// compile-flags: -O +// only-64bit (because the LLVM type of i64 for usize shows up) +// ignore-debug: all the `debug_assert`s get in the way + +#![crate_type = "lib"] +#![feature(array_repeat)] +#![feature(array_resize)] + +// CHECK-LABEL: @array_repeat_byte +#[no_mangle] +pub fn array_repeat_byte() -> [u8; 1234] { + // CHECK-NEXT: start: + // CHECK-NEXT: %1 = getelementptr + // CHECK-NEXT: tail call void @llvm.memset.p0i8.i64(i8* noundef nonnull align 1 dereferenceable(1234) %1, i8 42, i64 1234, i1 false) + // CHECK-NEXT: ret void + std::array::repeat(42) +} + +// CHECK-LABEL: @array_resize_byte_shrink +#[no_mangle] +pub fn array_resize_byte_shrink(a: [u8; 300]) -> [u8; 100] { + // CHECK-NOT: @llvm.memset + // CHECK: call void @llvm.memcpy.p0i8.p0i8.i64 + // CHECK-SAME: dereferenceable(100) + // CHECK-SAME: i64 100 + // CHECK-NOT: @llvm.memset + a.resize(42) +} + +// CHECK-LABEL: @array_resize_byte_grow +#[no_mangle] +pub fn array_resize_byte_grow(a: [u8; 100]) -> [u8; 300] { + // CHECK: call void @llvm.memcpy.p0i8.p0i8.i64 + // CHECK-SAME: dereferenceable(100) + // CHECK-SAME: i64 100 + // CHECK: call void @llvm.memset.p0i8.i64 + // CHECK-SAME: dereferenceable(200) + // CHECK-SAME: i8 42, i64 200 + a.resize(42) +} + +// CHECK-LABEL: @array_resize_with_byte_grow +#[no_mangle] +pub fn array_resize_with_byte_grow(a: [u8; 100]) -> [u8; 300] { + // CHECK: call void @llvm.memcpy.p0i8.p0i8.i64 + // CHECK-SAME: dereferenceable(100) + // CHECK-SAME: i64 100 + // CHECK: call void @llvm.memset.p0i8.i64 + // CHECK-SAME: dereferenceable(200) + // CHECK-SAME: i8 42, i64 200 + a.resize_with(|| 42) +} + +// CHECK-LABEL: @array_resize_string_moves_value +#[no_mangle] +pub fn array_resize_string_moves_value(a: [String; 1], b: String) -> [String; 2] { + // CHECK-NOT: __rust_dealloc + // CHECK: call void @llvm.memcpy.p0i8.p0i8.i64 + // CHECK-SAME: dereferenceable(24) + // CHECK-SAME: i64 24 + // CHECK-NOT: __rust_dealloc + // CHECK: call void @llvm.memcpy.p0i8.p0i8.i64 + // CHECK-SAME: dereferenceable(24) + // CHECK-SAME: i64 24 + // CHECK-NOT: __rust_dealloc + a.resize(b) +}