hybrid_array/
from_fn.rs

1//! Support for constructing arrays using a provided generator function.
2
3use crate::{Array, ArraySize};
4use core::{
5    convert::Infallible,
6    mem::{self, MaybeUninit},
7    ptr,
8};
9
10impl<T, U> Array<T, U>
11where
12    U: ArraySize,
13{
14    /// Create array where each array element `T` is returned by the `f` call.
15    #[inline]
16    pub fn from_fn(mut f: impl FnMut(usize) -> T) -> Self {
17        let Ok(ret) = Self::try_from_fn::<Infallible>(|n| Ok(f(n)));
18        ret
19    }
20
21    /// Create array fallibly where each array element `T` is returned by the `f` call, or return
22    /// an error if any are encountered.
23    ///
24    /// # Errors
25    ///
26    /// Propagates the `E` type returned from the provided `F` in the event of error.
27    pub fn try_from_fn<E>(f: impl FnMut(usize) -> Result<T, E>) -> Result<Self, E> {
28        let mut array = Array::<MaybeUninit<T>, U>::uninit();
29        try_from_fn_erased(array.0.as_mut(), f)?;
30
31        // SAFETY: if we got here, every element of the array was initialized
32        Ok(unsafe { array.assume_init() })
33    }
34}
35
36/// Fills a `MaybeUninit` slice using the given fallible generator function.
37///
38/// Using a slice avoids monomorphizing for each array size.
39#[inline]
40fn try_from_fn_erased<T, E, F>(buffer: &mut [MaybeUninit<T>], mut f: F) -> Result<(), E>
41where
42    F: FnMut(usize) -> Result<T, E>,
43{
44    let mut guard = Guard {
45        array_mut: buffer,
46        initialized: 0,
47    };
48
49    while guard.initialized < guard.array_mut.len() {
50        let item = f(guard.initialized)?;
51
52        // SAFETY: the loop's condition ensures we won't push too many items
53        unsafe { guard.push_unchecked(item) };
54    }
55
56    mem::forget(guard);
57    Ok(())
58}
59
60/// Drop guard which tracks the total number of initialized items, and handles dropping them in
61/// the event a panic occurs.
62///
63/// Use `mem::forget` when the array has been fully constructed.
64struct Guard<'a, T> {
65    /// Array being constructed.
66    array_mut: &'a mut [MaybeUninit<T>],
67
68    /// Number of items in the array which have been initialized.
69    initialized: usize,
70}
71
72impl<T> Guard<'_, T> {
73    /// Push an item onto the guard, writing to its `MaybeUninit` slot and incrementing the
74    /// counter of the number of initialized items.
75    ///
76    /// # Safety
77    ///
78    /// This can only be called n-times for as many elements are in the slice.
79    #[inline]
80    pub unsafe fn push_unchecked(&mut self, item: T) {
81        // SAFETY: the `initialized` counter tracks the number of initialized items, so as long as
82        // this is called the correct number of times for the array size writes will always be
83        // in-bounds and to an uninitialized slot in the array.
84        unsafe {
85            self.array_mut
86                .get_unchecked_mut(self.initialized)
87                .write(item);
88            self.initialized = self.initialized.saturating_add(1);
89        }
90    }
91}
92
93impl<T> Drop for Guard<'_, T> {
94    fn drop(&mut self) {
95        debug_assert!(self.initialized <= self.array_mut.len());
96
97        // SAFETY: the loop only iterates over initialized items
98        unsafe {
99            let p: *mut T = self.array_mut.as_mut_ptr().cast();
100
101            for i in 0..self.initialized {
102                ptr::drop_in_place(p.add(i));
103            }
104        }
105    }
106}
hybrid_array/from_fn.rs

hybrid_array/
from_fn.rs
