@@ -20,6 +20,69 @@ mod iter;
2020#[ stable( feature = "array_value_iter" , since = "1.51.0" ) ]
2121pub use iter:: IntoIter ;
2222
23+ /// Creates an array `[T; N]` where each array element `T` is returned by the `cb` call.
24+ ///
25+ /// # Arguments
26+ ///
27+ /// * `cb`: Callback where the passed argument is the current array index.
28+ ///
29+ /// # Example
30+ ///
31+ /// ```rust
32+ /// #![feature(array_from_fn)]
33+ ///
34+ /// let array = core::array::from_fn(|i| i);
35+ /// assert_eq!(array, [0, 1, 2, 3, 4]);
36+ /// ```
37+ #[ inline]
38+ #[ unstable( feature = "array_from_fn" , issue = "89379" ) ]
39+ pub fn from_fn < F , T , const N : usize > ( mut cb : F ) -> [ T ; N ]
40+ where
41+ F : FnMut ( usize ) -> T ,
42+ {
43+ let mut idx = 0 ;
44+ [ ( ) ; N ] . map ( |_| {
45+ let res = cb ( idx) ;
46+ idx += 1 ;
47+ res
48+ } )
49+ }
50+
51+ /// Creates an array `[T; N]` where each fallible array element `T` is returned by the `cb` call.
52+ /// Unlike `core::array::from_fn`, where the element creation can't fail, this version will return an error
53+ /// if any element creation was unsuccessful.
54+ ///
55+ /// # Arguments
56+ ///
57+ /// * `cb`: Callback where the passed argument is the current array index.
58+ ///
59+ /// # Example
60+ ///
61+ /// ```rust
62+ /// #![feature(array_from_fn)]
63+ ///
64+ /// #[derive(Debug, PartialEq)]
65+ /// enum SomeError {
66+ /// Foo,
67+ /// }
68+ ///
69+ /// let array = core::array::try_from_fn(|i| Ok::<_, SomeError>(i));
70+ /// assert_eq!(array, Ok([0, 1, 2, 3, 4]));
71+ ///
72+ /// let another_array = core::array::try_from_fn::<SomeError, _, (), 2>(|_| Err(SomeError::Foo));
73+ /// assert_eq!(another_array, Err(SomeError::Foo));
74+ /// ```
75+ #[ inline]
76+ #[ unstable( feature = "array_from_fn" , issue = "89379" ) ]
77+ pub fn try_from_fn < E , F , T , const N : usize > ( cb : F ) -> Result < [ T ; N ] , E >
78+ where
79+ F : FnMut ( usize ) -> Result < T , E > ,
80+ {
81+ // SAFETY: we know for certain that this iterator will yield exactly `N`
82+ // items.
83+ unsafe { collect_into_array_rslt_unchecked ( & mut ( 0 ..N ) . map ( cb) ) }
84+ }
85+
2386/// Converts a reference to `T` into a reference to an array of length 1 (without copying).
2487#[ stable( feature = "array_from_ref" , since = "1.53.0" ) ]
2588pub fn from_ref < T > ( s : & T ) -> & [ T ; 1 ] {
@@ -448,13 +511,15 @@ impl<T, const N: usize> [T; N] {
448511///
449512/// It is up to the caller to guarantee that `iter` yields at least `N` items.
450513/// Violating this condition causes undefined behavior.
451- unsafe fn collect_into_array_unchecked < I , const N : usize > ( iter : & mut I ) -> [ I :: Item ; N ]
514+ unsafe fn collect_into_array_rslt_unchecked < E , I , T , const N : usize > (
515+ iter : & mut I ,
516+ ) -> Result < [ T ; N ] , E >
452517where
453518 // Note: `TrustedLen` here is somewhat of an experiment. This is just an
454519 // internal function, so feel free to remove if this bound turns out to be a
455520 // bad idea. In that case, remember to also remove the lower bound
456521 // `debug_assert!` below!
457- I : Iterator + TrustedLen ,
522+ I : Iterator < Item = Result < T , E > > + TrustedLen ,
458523{
459524 debug_assert ! ( N <= iter. size_hint( ) . 1 . unwrap_or( usize :: MAX ) ) ;
460525 debug_assert ! ( N <= iter. size_hint( ) . 0 ) ;
@@ -463,6 +528,21 @@ where
463528 unsafe { collect_into_array ( iter) . unwrap_unchecked ( ) }
464529}
465530
531+ // Infallible version of `collect_into_array_rslt_unchecked`.
532+ unsafe fn collect_into_array_unchecked < I , const N : usize > ( iter : & mut I ) -> [ I :: Item ; N ]
533+ where
534+ I : Iterator + TrustedLen ,
535+ {
536+ let mut map = iter. map ( Ok :: < _ , Infallible > ) ;
537+
538+ // SAFETY: The same safety considerations w.r.t. the iterator length
539+ // apply for `collect_into_array_rslt_unchecked` as for
540+ // `collect_into_array_unchecked`
541+ match unsafe { collect_into_array_rslt_unchecked ( & mut map) } {
542+ Ok ( array) => array,
543+ }
544+ }
545+
466546/// Pulls `N` items from `iter` and returns them as an array. If the iterator
467547/// yields fewer than `N` items, `None` is returned and all already yielded
468548/// items are dropped.
@@ -473,43 +553,49 @@ where
473553///
474554/// If `iter.next()` panicks, all items already yielded by the iterator are
475555/// dropped.
476- fn collect_into_array < I , const N : usize > ( iter : & mut I ) -> Option < [ I :: Item ; N ] >
556+ fn collect_into_array < E , I , T , const N : usize > ( iter : & mut I ) -> Option < Result < [ T ; N ] , E > >
477557where
478- I : Iterator ,
558+ I : Iterator < Item = Result < T , E > > ,
479559{
480560 if N == 0 {
481561 // SAFETY: An empty array is always inhabited and has no validity invariants.
482- return unsafe { Some ( mem:: zeroed ( ) ) } ;
562+ return unsafe { Some ( Ok ( mem:: zeroed ( ) ) ) } ;
483563 }
484564
485- struct Guard < T , const N : usize > {
486- ptr : * mut T ,
565+ struct Guard < ' a , T , const N : usize > {
566+ array_mut : & ' a mut [ MaybeUninit < T > ; N ] ,
487567 initialized : usize ,
488568 }
489569
490- impl < T , const N : usize > Drop for Guard < T , N > {
570+ impl < T , const N : usize > Drop for Guard < ' _ , T , N > {
491571 fn drop ( & mut self ) {
492572 debug_assert ! ( self . initialized <= N ) ;
493573
494- let initialized_part = crate :: ptr:: slice_from_raw_parts_mut ( self . ptr , self . initialized ) ;
495-
496- // SAFETY: this raw slice will contain only initialized objects.
574+ // SAFETY: this slice will contain only initialized objects.
497575 unsafe {
498- crate :: ptr:: drop_in_place ( initialized_part) ;
576+ crate :: ptr:: drop_in_place ( MaybeUninit :: slice_assume_init_mut (
577+ & mut self . array_mut . get_unchecked_mut ( ..self . initialized ) ,
578+ ) ) ;
499579 }
500580 }
501581 }
502582
503583 let mut array = MaybeUninit :: uninit_array :: < N > ( ) ;
504- let mut guard: Guard < _ , N > =
505- Guard { ptr : MaybeUninit :: slice_as_mut_ptr ( & mut array) , initialized : 0 } ;
584+ let mut guard = Guard { array_mut : & mut array, initialized : 0 } ;
585+
586+ while let Some ( item_rslt) = iter. next ( ) {
587+ let item = match item_rslt {
588+ Err ( err) => {
589+ return Some ( Err ( err) ) ;
590+ }
591+ Ok ( elem) => elem,
592+ } ;
506593
507- while let Some ( item) = iter. next ( ) {
508594 // SAFETY: `guard.initialized` starts at 0, is increased by one in the
509595 // loop and the loop is aborted once it reaches N (which is
510596 // `array.len()`).
511597 unsafe {
512- array . get_unchecked_mut ( guard. initialized ) . write ( item) ;
598+ guard . array_mut . get_unchecked_mut ( guard. initialized ) . write ( item) ;
513599 }
514600 guard. initialized += 1 ;
515601
@@ -520,7 +606,7 @@ where
520606 // SAFETY: the condition above asserts that all elements are
521607 // initialized.
522608 let out = unsafe { MaybeUninit :: array_assume_init ( array) } ;
523- return Some ( out) ;
609+ return Some ( Ok ( out) ) ;
524610 }
525611 }
526612
0 commit comments