Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add core::array::from_fn #83576

Closed
wants to merge 1 commit into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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]);
```
CryZe committed Mar 27, 2021
commit 3a4a874185d9ad0078e6f71f698f824e1b0b9084
82 changes: 63 additions & 19 deletions library/core/src/array/mod.rs
Original file line number Diff line number Diff line change
@@ -532,6 +532,68 @@ impl<T, const N: usize> [T; N] {
}
}

struct Guard<T, const N: usize> {
ptr: *mut T,
initialized: usize,
}

impl<T, const N: usize> Drop for Guard<T, N> {
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<T, F, const N: usize>(mut f: F) -> [T; N]
where
F: FnMut(usize) -> T,
{
let mut array = MaybeUninit::uninit_array::<N>();
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<I, const N: usize>(iter: &mut I) -> Option<[I::Item; N]>
where
@@ -579,24 +641,6 @@ where
return unsafe { Some(mem::zeroed()) };
}

struct Guard<T, const N: usize> {
ptr: *mut T,
initialized: usize,
}

impl<T, const N: usize> Drop for Guard<T, N> {
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::<N>();
let mut guard: Guard<_, N> =
Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 };
9 changes: 9 additions & 0 deletions library/core/tests/array.rs
Original file line number Diff line number Diff line change
@@ -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];
1 change: 1 addition & 0 deletions library/core/tests/lib.rs
Original file line number Diff line number Diff line change
@@ -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)]