diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 8f52985d1df71..fd8827d0bcc89 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -532,6 +532,68 @@ impl [T; N] { } } +struct Guard { + ptr: *mut T, + initialized: usize, +} + +impl Drop for Guard { + fn drop(&mut self) { + debug_assert!(self.initialized <= N); + + let initialized_part = crate::ptr::slice_from_raw_parts_mut(self.ptr, self.initialized); + + // SAFETY: this raw slice will contain only initialized objects. + unsafe { + crate::ptr::drop_in_place(initialized_part); + } + } +} + +/// Creates a new array where each element is initialized by calling the +/// provided closure `F: FnMut(usize) -> T` with the index of each element. +/// +/// This allows initializing an array with elements that are not `Copy` or +/// `const` and where writing out the entire array literal as opposed to +/// `[expr; N]` is too verbose or not possible. +/// +/// # Example +/// +/// ``` +/// #![feature(array_from_fn)] +/// +/// let array = std::array::from_fn(|index| 2 * index); +/// assert_eq!(array, [0, 2, 4, 6]); +/// ``` +#[unstable(feature = "array_from_fn", issue = "none")] +pub fn from_fn(mut f: F) -> [T; N] +where + F: FnMut(usize) -> T, +{ + let mut array = MaybeUninit::uninit_array::(); + let mut guard: Guard<_, N> = + Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 }; + + while guard.initialized < N { + let item = f(guard.initialized); + + // SAFETY: `guard.initialized` starts at 0, is increased by one in the + // loop and the loop is aborted once it reaches N (which is + // `array.len()`). + unsafe { + array.get_unchecked_mut(guard.initialized).write(item); + } + + guard.initialized += 1; + } + + mem::forget(guard); + + // SAFETY: the condition above asserts that all elements are + // initialized. + unsafe { MaybeUninit::array_assume_init(array) } +} + /// Pulls `N` items from `iter` and returns them as an array. If the iterator /// yields fewer than `N` items, this function exhibits undefined behavior. /// @@ -568,7 +630,7 @@ where /// `next` at most `N` times, the iterator can still be used afterwards to /// retrieve the remaining items. /// -/// If `iter.next()` panicks, all items already yielded by the iterator are +/// If `iter.next()` panics, all items already yielded by the iterator are /// dropped. fn collect_into_array(iter: &mut I) -> Option<[I::Item; N]> where @@ -579,24 +641,6 @@ where return unsafe { Some(mem::zeroed()) }; } - struct Guard { - ptr: *mut T, - initialized: usize, - } - - impl Drop for Guard { - fn drop(&mut self) { - debug_assert!(self.initialized <= N); - - let initialized_part = crate::ptr::slice_from_raw_parts_mut(self.ptr, self.initialized); - - // SAFETY: this raw slice will contain only initialized objects. - unsafe { - crate::ptr::drop_in_place(initialized_part); - } - } - } - let mut array = MaybeUninit::uninit_array::(); let mut guard: Guard<_, N> = Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 }; diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs index 89c2a969c28bb..0452876bcee85 100644 --- a/library/core/tests/array.rs +++ b/library/core/tests/array.rs @@ -306,6 +306,15 @@ fn empty_array_is_always_default() { let _arr = <[DoesNotImplDefault; 0]>::default(); } +#[test] +fn array_from_fn() { + let a = array::from_fn(|i| 2 * i + 1); + assert_eq!(a, [1, 3, 5, 7]); + + let b = array::from_fn(|i| i.to_string()); + assert_eq!(b, [String::from("0"), String::from("1")]); +} + #[test] fn array_map() { let a = [1, 2, 3]; diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 997e618cc511f..e862e3d220c1d 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -1,5 +1,6 @@ #![feature(alloc_layout_extra)] #![feature(array_chunks)] +#![feature(array_from_fn)] #![feature(array_from_ref)] #![feature(array_methods)] #![feature(array_map)]