diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 811850af3678d..79bb48611b48f 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -11,7 +11,7 @@ use crate::fmt; use crate::hash::{self, Hash}; use crate::iter::TrustedLen; use crate::mem::{self, MaybeUninit}; -use crate::ops::{Index, IndexMut}; +use crate::ops::{ControlFlow, Index, IndexMut, Try}; use crate::slice::{Iter, IterMut}; mod equality; @@ -24,37 +24,37 @@ pub use iter::IntoIter; /// /// # Arguments /// -/// * `cb`: Callback where the passed argument is the current array index. +/// * `cb`: Callback that must provide the array element /// /// # Example /// /// ```rust /// #![feature(array_from_fn)] /// -/// let array = core::array::from_fn(|i| i); +/// let mut idx: usize = 0; +/// let array = core::array::from_fn(|| { +/// let rslt = idx; +/// idx += 1; +/// rslt +/// }); /// 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, + F: FnMut() -> T, { - let mut idx = 0; - [(); N].map(|_| { - let res = cb(idx); - idx += 1; - res - }) + [(); N].map(|_| cb()) } /// 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 +/// 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. +/// * `cb`: Callback that must provide the array element /// /// # Example /// @@ -66,21 +66,23 @@ where /// Foo, /// } /// -/// let array = core::array::try_from_fn(|i| Ok::<_, SomeError>(i)); -/// assert_eq!(array, Ok([0, 1, 2, 3, 4])); +/// let array: Result<[i32; 5], SomeError> = core::array::try_from_fn(|| Ok(1)); +/// assert_eq!(array, Ok([1, 1, 1, 1, 1])); /// -/// let another_array = core::array::try_from_fn::(|_| Err(SomeError::Foo)); +/// let another_array: Result<[(); 2], SomeError> = 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> +pub fn try_from_fn(mut cb: F) -> AT where - F: FnMut(usize) -> Result, + AT: Try, + ET: Try, + F: FnMut() -> ET, { // SAFETY: we know for certain that this iterator will yield exactly `N` // items. - unsafe { collect_into_array_rslt_unchecked(&mut (0..N).map(cb)) } + 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). @@ -592,15 +594,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_rslt_unchecked( - iter: &mut I, -) -> Result<[T; N], E> +unsafe fn collect_into_array_rslt_unchecked(iter: &mut I) -> AT where + AT: Try, + ET: Try, // 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); @@ -609,7 +611,7 @@ where unsafe { collect_into_array(iter).unwrap_unchecked() } } -// Infallible version of `collect_into_array_rslt_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, @@ -634,13 +636,15 @@ where /// /// If `iter.next()` panicks, all items already yielded by the iterator are /// dropped. -fn collect_into_array(iter: &mut I) -> Option> +fn collect_into_array(iter: &mut I) -> Option where - I: Iterator>, + AT: Try, + ET: Try, + I: Iterator, { if N == 0 { // SAFETY: An empty array is always inhabited and has no validity invariants. - return unsafe { Some(Ok(mem::zeroed())) }; + return unsafe { Some(AT::from_output(mem::zeroed())) }; } struct Guard<'a, T, const N: usize> { @@ -665,11 +669,11 @@ where 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)); + let item = match item_rslt.branch() { + ControlFlow::Break(res) => { + return Some(AT::from_residual(res)); } - Ok(elem) => elem, + ControlFlow::Continue(elem) => elem, }; // SAFETY: `guard.initialized` starts at 0, is increased by one in the @@ -687,7 +691,7 @@ where // SAFETY: the condition above asserts that all elements are // initialized. let out = unsafe { MaybeUninit::array_assume_init(array) }; - return Some(Ok(out)); + return Some(AT::from_output(out)); } } diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs index 7dc071b74235d..7b9811068b04f 100644 --- a/library/core/tests/array.rs +++ b/library/core/tests/array.rs @@ -363,8 +363,10 @@ fn cell_allows_array_cycle() { #[test] fn array_from_fn() { - let array = core::array::from_fn(|idx| idx); - assert_eq!(array, [0, 1, 2, 3, 4]); + #[derive(Debug, PartialEq)] + struct Foo; + let array = core::array::from_fn(|| Foo); + assert_eq!(array, [Foo, Foo]); } #[test] @@ -374,10 +376,11 @@ fn array_try_from_fn() { Foo, } - let array = core::array::try_from_fn(|i| Ok::<_, SomeError>(i)); - assert_eq!(array, Ok([0, 1, 2, 3, 4])); + let array: Result<[i32; 5], SomeError> = core::array::try_from_fn(|| Ok(1)); + assert_eq!(array, Ok([1, 1, 1, 1, 1])); - let another_array = core::array::try_from_fn::(|_| Err(SomeError::Foo)); + let another_array: Result<[(); 2], SomeError> = + core::array::try_from_fn(|| Err(SomeError::Foo)); assert_eq!(another_array, Err(SomeError::Foo)); } @@ -394,10 +397,12 @@ fn array_try_from_fn_drops_inserted_elements_on_err() { } let _ = catch_unwind_silent(move || { - let _: Result<[CountDrop; 4], ()> = core::array::try_from_fn(|idx| { + let mut idx = 0; + let _: Result<[CountDrop; 4], ()> = core::array::try_from_fn(|| { if idx == 2 { return Err(()); } + idx += 1; Ok(CountDrop) }); }); @@ -418,10 +423,12 @@ fn array_try_from_fn_drops_inserted_elements_on_panic() { } let _ = catch_unwind_silent(move || { - let _: Result<[CountDrop; 4], ()> = core::array::try_from_fn(|idx| { + let mut idx = 0; + let _: Result<[CountDrop; 4], ()> = core::array::try_from_fn(|| { if idx == 2 { panic!("peek a boo"); } + idx += 1; Ok(CountDrop) }); });