1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

//! Support for constructing arrays using a provided generator function.

use crate::{Array, ArraySize};
use core::{
    mem::{self, MaybeUninit},
    ptr,
};

/// Construct an array type from the given generator function.
pub trait FromFn<T>: Sized {
    /// Create array using the given generator function for each element.
    fn from_fn<F>(cb: F) -> Self
    where
        F: FnMut(usize) -> T;

    /// Create an array using the given generator function for each element, returning any errors
    /// which are encountered in the given generator.
    fn try_from_fn<E, F>(cb: F) -> Result<Self, E>
    where
        F: FnMut(usize) -> Result<T, E>;
}

impl<T, U> FromFn<T> for Array<T, U>
where
    U: ArraySize,
{
    #[inline]
    fn from_fn<F>(cb: F) -> Self
    where
        F: FnMut(usize) -> T,
    {
        Array::from_fn(cb)
    }

    #[inline]
    fn try_from_fn<E, F>(cb: F) -> Result<Self, E>
    where
        F: FnMut(usize) -> Result<T, E>,
    {
        Array::try_from_fn(cb)
    }
}

impl<T, const N: usize> FromFn<T> for [T; N] {
    #[inline]
    fn from_fn<F>(cb: F) -> Self
    where
        F: FnMut(usize) -> T,
    {
        core::array::from_fn(cb)
    }

    // TODO(tarcieri): use `array::try_from_fn` when stable
    fn try_from_fn<E, F>(cb: F) -> Result<Self, E>
    where
        F: FnMut(usize) -> Result<T, E>,
    {
        // SAFETY: an array of `MaybeUninit`s is always valid.
        let mut array: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init() };
        try_from_fn_erased(&mut array, cb)?;

        // TODO(tarcieri): use `MaybeUninit::array_assume_init` when stable
        // SAFETY: if we got here, every element of the array was initialized
        Ok(unsafe { ptr::read(array.as_ptr().cast()) })
    }
}

/// Fills a `MaybeUninit` slice using the given fallible generator function.
///
/// Using a slice avoids monomorphizing for each array size.
#[inline]
fn try_from_fn_erased<T, E, F>(buffer: &mut [MaybeUninit<T>], mut cb: F) -> Result<(), E>
where
    F: FnMut(usize) -> Result<T, E>,
{
    let mut guard = Guard {
        array_mut: buffer,
        initialized: 0,
    };

    while guard.initialized < guard.array_mut.len() {
        let item = cb(guard.initialized)?;

        // SAFETY: the loop's condition ensures we won't push too many items
        unsafe { guard.push_unchecked(item) };
    }

    mem::forget(guard);
    Ok(())
}

/// Drop guard which tracks the total number of initialized items, and handles dropping them in
/// the event a panic occurs.
///
/// Use `mem::forget` when the array has been fully constructed.
struct Guard<'a, T> {
    /// Array being constructed.
    array_mut: &'a mut [MaybeUninit<T>],

    /// Number of items in the array which have been initialized.
    initialized: usize,
}

impl<T> Guard<'_, T> {
    /// Push an item onto the guard, writing to its `MaybeUninit` slot and incrementing the
    /// counter of the number of initialized items.
    ///
    /// # Safety
    ///
    /// This can only be called n-times for as many elements are in the slice.
    #[inline]
    pub unsafe fn push_unchecked(&mut self, item: T) {
        // SAFETY: the `initialized` counter tracks the number of initialized items, so as long as
        // this is called the correct number of times for the array size writes will always be
        // in-bounds and to an uninitialized slot in the array.
        unsafe {
            self.array_mut
                .get_unchecked_mut(self.initialized)
                .write(item);
            self.initialized = self.initialized.saturating_add(1);
        }
    }
}

impl<T> Drop for Guard<'_, T> {
    fn drop(&mut self) {
        debug_assert!(self.initialized <= self.array_mut.len());

        // SAFETY: the loop only iterates over initialized items
        unsafe {
            let p: *mut T = self.array_mut.as_mut_ptr().cast();

            for i in 0..self.initialized {
                ptr::drop_in_place(p.add(i));
            }
        }
    }
}