Skip to content

Commit 3a4a874

Browse files
committed
Add core::array::from_fn
This adds `core::array::from_fn` for initializing an array from a closure. The API is meant to closely resemble `core::iter::from_fn` with the difference that the array dictates the amount of elements produced as opposed to the function. This should also cover a lot of cases where one might want to reach for collecting into an array, but without all the associated problems that this would have, such as providing too few or too many elements. ```rust let array = std::array::from_fn(|index| 2 * index); assert_eq!(array, [0, 2, 4, 6]); ```
1 parent aef1140 commit 3a4a874

File tree

3 files changed

+73
-19
lines changed

3 files changed

+73
-19
lines changed

library/core/src/array/mod.rs

+63-19
Original file line numberDiff line numberDiff line change
@@ -532,6 +532,68 @@ impl<T, const N: usize> [T; N] {
532532
}
533533
}
534534

535+
struct Guard<T, const N: usize> {
536+
ptr: *mut T,
537+
initialized: usize,
538+
}
539+
540+
impl<T, const N: usize> Drop for Guard<T, N> {
541+
fn drop(&mut self) {
542+
debug_assert!(self.initialized <= N);
543+
544+
let initialized_part = crate::ptr::slice_from_raw_parts_mut(self.ptr, self.initialized);
545+
546+
// SAFETY: this raw slice will contain only initialized objects.
547+
unsafe {
548+
crate::ptr::drop_in_place(initialized_part);
549+
}
550+
}
551+
}
552+
553+
/// Creates a new array where each element is initialized by calling the
554+
/// provided closure `F: FnMut(usize) -> T` with the index of each element.
555+
///
556+
/// This allows initializing an array with elements that are not `Copy` or
557+
/// `const` and where writing out the entire array literal as opposed to
558+
/// `[expr; N]` is too verbose or not possible.
559+
///
560+
/// # Example
561+
///
562+
/// ```
563+
/// #![feature(array_from_fn)]
564+
///
565+
/// let array = std::array::from_fn(|index| 2 * index);
566+
/// assert_eq!(array, [0, 2, 4, 6]);
567+
/// ```
568+
#[unstable(feature = "array_from_fn", issue = "none")]
569+
pub fn from_fn<T, F, const N: usize>(mut f: F) -> [T; N]
570+
where
571+
F: FnMut(usize) -> T,
572+
{
573+
let mut array = MaybeUninit::uninit_array::<N>();
574+
let mut guard: Guard<_, N> =
575+
Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 };
576+
577+
while guard.initialized < N {
578+
let item = f(guard.initialized);
579+
580+
// SAFETY: `guard.initialized` starts at 0, is increased by one in the
581+
// loop and the loop is aborted once it reaches N (which is
582+
// `array.len()`).
583+
unsafe {
584+
array.get_unchecked_mut(guard.initialized).write(item);
585+
}
586+
587+
guard.initialized += 1;
588+
}
589+
590+
mem::forget(guard);
591+
592+
// SAFETY: the condition above asserts that all elements are
593+
// initialized.
594+
unsafe { MaybeUninit::array_assume_init(array) }
595+
}
596+
535597
/// Pulls `N` items from `iter` and returns them as an array. If the iterator
536598
/// yields fewer than `N` items, this function exhibits undefined behavior.
537599
///
@@ -568,7 +630,7 @@ where
568630
/// `next` at most `N` times, the iterator can still be used afterwards to
569631
/// retrieve the remaining items.
570632
///
571-
/// If `iter.next()` panicks, all items already yielded by the iterator are
633+
/// If `iter.next()` panics, all items already yielded by the iterator are
572634
/// dropped.
573635
fn collect_into_array<I, const N: usize>(iter: &mut I) -> Option<[I::Item; N]>
574636
where
@@ -579,24 +641,6 @@ where
579641
return unsafe { Some(mem::zeroed()) };
580642
}
581643

582-
struct Guard<T, const N: usize> {
583-
ptr: *mut T,
584-
initialized: usize,
585-
}
586-
587-
impl<T, const N: usize> Drop for Guard<T, N> {
588-
fn drop(&mut self) {
589-
debug_assert!(self.initialized <= N);
590-
591-
let initialized_part = crate::ptr::slice_from_raw_parts_mut(self.ptr, self.initialized);
592-
593-
// SAFETY: this raw slice will contain only initialized objects.
594-
unsafe {
595-
crate::ptr::drop_in_place(initialized_part);
596-
}
597-
}
598-
}
599-
600644
let mut array = MaybeUninit::uninit_array::<N>();
601645
let mut guard: Guard<_, N> =
602646
Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 };

library/core/tests/array.rs

+9
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,15 @@ fn empty_array_is_always_default() {
306306
let _arr = <[DoesNotImplDefault; 0]>::default();
307307
}
308308

309+
#[test]
310+
fn array_from_fn() {
311+
let a = array::from_fn(|i| 2 * i + 1);
312+
assert_eq!(a, [1, 3, 5, 7]);
313+
314+
let b = array::from_fn(|i| i.to_string());
315+
assert_eq!(b, [String::from("0"), String::from("1")]);
316+
}
317+
309318
#[test]
310319
fn array_map() {
311320
let a = [1, 2, 3];

library/core/tests/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#![feature(alloc_layout_extra)]
22
#![feature(array_chunks)]
3+
#![feature(array_from_fn)]
34
#![feature(array_from_ref)]
45
#![feature(array_methods)]
56
#![feature(array_map)]

0 commit comments

Comments
 (0)