diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 09329247f9492..8d5c0510404fb 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -20,6 +20,69 @@ mod iter; #[stable(feature = "array_value_iter", since = "1.51.0")] pub use iter::IntoIter; +/// Creates an array `[T; N]` where each array element `T` is returned by the `cb` call. +/// +/// # Arguments +/// +/// * `cb`: Callback where the passed argument is the current array index. +/// +/// # Example +/// +/// ```rust +/// #![feature(array_from_fn)] +/// +/// let array = core::array::from_fn(|i| i); +/// assert_eq!(array, [0, 1, 2, 3, 4]); +/// ``` +#[inline] +#[unstable(feature = "array_from_fn", issue = "89379")] +pub fn from_fn(mut cb: F) -> [T; N] +where + F: FnMut(usize) -> T, +{ + let mut idx = 0; + [(); N].map(|_| { + let res = cb(idx); + idx += 1; + res + }) +} + +/// Creates an array `[T; N]` where each fallible array element `T` is returned by the `cb` call. +/// Unlike `core::array::from_fn`, where the element creation can't fail, this version will return an error +/// if any element creation was unsuccessful. +/// +/// # Arguments +/// +/// * `cb`: Callback where the passed argument is the current array index. +/// +/// # Example +/// +/// ```rust +/// #![feature(array_from_fn)] +/// +/// #[derive(Debug, PartialEq)] +/// enum SomeError { +/// Foo, +/// } +/// +/// let array = core::array::try_from_fn(|i| Ok::<_, SomeError>(i)); +/// assert_eq!(array, Ok([0, 1, 2, 3, 4])); +/// +/// let another_array = core::array::try_from_fn::(|_| Err(SomeError::Foo)); +/// assert_eq!(another_array, Err(SomeError::Foo)); +/// ``` +#[inline] +#[unstable(feature = "array_from_fn", issue = "89379")] +pub fn try_from_fn(cb: F) -> Result<[T; N], E> +where + F: FnMut(usize) -> Result, +{ + // SAFETY: we know for certain that this iterator will yield exactly `N` + // items. + unsafe { collect_into_array_rslt_unchecked(&mut (0..N).map(cb)) } +} + /// 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")] pub fn from_ref(s: &T) -> &[T; 1] { @@ -448,13 +511,15 @@ impl [T; N] { /// /// It is up to the caller to guarantee that `iter` yields at least `N` items. /// Violating this condition causes undefined behavior. -unsafe fn collect_into_array_unchecked(iter: &mut I) -> [I::Item; N] +unsafe fn collect_into_array_rslt_unchecked( + iter: &mut I, +) -> Result<[T; N], E> where // Note: `TrustedLen` here is somewhat of an experiment. This is just an // internal function, so feel free to remove if this bound turns out to be a // bad idea. In that case, remember to also remove the lower bound // `debug_assert!` below! - I: Iterator + TrustedLen, + I: Iterator> + TrustedLen, { debug_assert!(N <= iter.size_hint().1.unwrap_or(usize::MAX)); debug_assert!(N <= iter.size_hint().0); @@ -463,6 +528,21 @@ where unsafe { collect_into_array(iter).unwrap_unchecked() } } +// Infallible version of `collect_into_array_rslt_unchecked`. +unsafe fn collect_into_array_unchecked(iter: &mut I) -> [I::Item; N] +where + I: Iterator + TrustedLen, +{ + let mut map = iter.map(Ok::<_, Infallible>); + + // SAFETY: The same safety considerations w.r.t. the iterator length + // apply for `collect_into_array_rslt_unchecked` as for + // `collect_into_array_unchecked` + match unsafe { collect_into_array_rslt_unchecked(&mut map) } { + Ok(array) => array, + } +} + /// Pulls `N` items from `iter` and returns them as an array. If the iterator /// yields fewer than `N` items, `None` is returned and all already yielded /// items are dropped. @@ -473,43 +553,49 @@ where /// /// If `iter.next()` panicks, all items already yielded by the iterator are /// dropped. -fn collect_into_array(iter: &mut I) -> Option<[I::Item; N]> +fn collect_into_array(iter: &mut I) -> Option> where - I: Iterator, + I: Iterator>, { if N == 0 { // SAFETY: An empty array is always inhabited and has no validity invariants. - return unsafe { Some(mem::zeroed()) }; + return unsafe { Some(Ok(mem::zeroed())) }; } - struct Guard { - ptr: *mut T, + struct Guard<'a, T, const N: usize> { + array_mut: &'a mut [MaybeUninit; N], initialized: usize, } - impl Drop for Guard { + impl 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. + // SAFETY: this slice will contain only initialized objects. unsafe { - crate::ptr::drop_in_place(initialized_part); + 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<_, N> = - Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 }; + let mut guard = Guard { array_mut: &mut array, initialized: 0 }; + + while let Some(item_rslt) = iter.next() { + let item = match item_rslt { + Err(err) => { + return Some(Err(err)); + } + Ok(elem) => elem, + }; - while let Some(item) = iter.next() { // 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.array_mut.get_unchecked_mut(guard.initialized).write(item); } guard.initialized += 1; @@ -520,7 +606,7 @@ where // SAFETY: the condition above asserts that all elements are // initialized. let out = unsafe { MaybeUninit::array_assume_init(array) }; - return Some(out); + return Some(Ok(out)); } } diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs index 0ae625bdb68c6..b3af1328c90d4 100644 --- a/library/core/tests/array.rs +++ b/library/core/tests/array.rs @@ -1,5 +1,6 @@ use core::array; use core::convert::TryFrom; +use core::sync::atomic::{AtomicUsize, Ordering}; #[test] fn array_from_ref() { @@ -303,8 +304,6 @@ fn array_map() { #[test] #[should_panic(expected = "test succeeded")] fn array_map_drop_safety() { - use core::sync::atomic::AtomicUsize; - use core::sync::atomic::Ordering; static DROPPED: AtomicUsize = AtomicUsize::new(0); struct DropCounter; impl Drop for DropCounter { @@ -356,3 +355,84 @@ fn cell_allows_array_cycle() { b3.a[0].set(Some(&b1)); b3.a[1].set(Some(&b2)); } + +#[test] +fn array_from_fn() { + let array = core::array::from_fn(|idx| idx); + assert_eq!(array, [0, 1, 2, 3, 4]); +} + +#[test] +fn array_try_from_fn() { + #[derive(Debug, PartialEq)] + enum SomeError { + Foo, + } + + let array = core::array::try_from_fn(|i| Ok::<_, SomeError>(i)); + assert_eq!(array, Ok([0, 1, 2, 3, 4])); + + let another_array = core::array::try_from_fn::(|_| Err(SomeError::Foo)); + assert_eq!(another_array, Err(SomeError::Foo)); +} + +#[cfg(not(panic = "abort"))] +#[test] +fn array_try_from_fn_drops_inserted_elements_on_err() { + static DROP_COUNTER: AtomicUsize = AtomicUsize::new(0); + + struct CountDrop; + impl Drop for CountDrop { + fn drop(&mut self) { + DROP_COUNTER.fetch_add(1, Ordering::SeqCst); + } + } + + let _ = catch_unwind_silent(move || { + let _: Result<[CountDrop; 4], ()> = core::array::try_from_fn(|idx| { + if idx == 2 { + return Err(()); + } + Ok(CountDrop) + }); + }); + + assert_eq!(DROP_COUNTER.load(Ordering::SeqCst), 2); +} + +#[cfg(not(panic = "abort"))] +#[test] +fn array_try_from_fn_drops_inserted_elements_on_panic() { + static DROP_COUNTER: AtomicUsize = AtomicUsize::new(0); + + struct CountDrop; + impl Drop for CountDrop { + fn drop(&mut self) { + DROP_COUNTER.fetch_add(1, Ordering::SeqCst); + } + } + + let _ = catch_unwind_silent(move || { + let _: Result<[CountDrop; 4], ()> = core::array::try_from_fn(|idx| { + if idx == 2 { + panic!("peek a boo"); + } + Ok(CountDrop) + }); + }); + + assert_eq!(DROP_COUNTER.load(Ordering::SeqCst), 2); +} + +#[cfg(not(panic = "abort"))] +// https://stackoverflow.com/a/59211505 +fn catch_unwind_silent(f: F) -> std::thread::Result +where + F: FnOnce() -> R + core::panic::UnwindSafe, +{ + let prev_hook = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); + let result = std::panic::catch_unwind(f); + std::panic::set_hook(prev_hook); + result +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 0673564dd5c4a..3608853dce4e0 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -26,6 +26,7 @@ #![feature(extern_types)] #![feature(flt2dec)] #![feature(fmt_internals)] +#![feature(array_from_fn)] #![feature(hashmap_internals)] #![feature(try_find)] #![feature(is_sorted)]