From 4be574e6c91e004c13e133ea3a441e13aa2439d3 Mon Sep 17 00:00:00 2001 From: Caio Date: Thu, 30 Sep 2021 07:49:32 -0300 Subject: Add 'core::array::from_fn' and 'core::array::try_from_fn' --- library/core/src/array/mod.rs | 98 +++++++++++++++++++++++++++++++++++++++---- 1 file changed, 91 insertions(+), 7 deletions(-) (limited to 'library/core/src/array') diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 8c33a43ab33..b823c33ee23 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,18 @@ 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(|el| Ok::<_, Infallible>(el)); + + // SAFETY: Valid array elements are covered by the fact that all passed values + // to `collect_into_array` are `Ok`. + unsafe { collect_into_array_rslt_unchecked(&mut map).unwrap_unchecked() } +} + /// 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,13 +550,13 @@ 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 { @@ -504,7 +581,14 @@ where let mut guard: Guard<_, N> = Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 }; - while let Some(item) = iter.next() { + while let Some(item_rslt) = iter.next() { + let item = match item_rslt { + Err(err) => { + return Some(Err(err)); + } + Ok(elem) => elem, + }; + // 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()`). @@ -520,7 +604,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)); } } -- cgit 1.4.1-3-g733a5