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

hybrid_array/
from_fn.rs
