wasmtime_internal_core/
array.rs

1//! Not yet stable array functions needed by fixed length list
2
3use core::mem::MaybeUninit;
4
5// See https://doc.rust-lang.org/core/array/fn.try_from_fn.html for the not yet stable original
6//
7/// Creates an array `[T; N]` where each fallible array element `T` is returned by the `cb` call.
8/// Unlike [`core::array::from_fn`], where the element creation can't fail, this version will return an error
9/// if any element creation was unsuccessful.
10///
11/// The return type of this function depends on the return type of the closure.
12/// If you return `Result<T, E>` from the closure, you'll get a `Result<[T; N], E>`.
13///
14/// Note: Unlike the unstable core implementation this function only supports a closure returning a Result.
15///
16/// # Arguments
17///
18/// * `cb`: Callback where the passed argument is the current array index.
19///
20/// # Example
21///
22/// ```rust
23/// use wasmtime_internal_core::array::array_try_from_fn;
24///
25/// let array: Result<[u8; 5], _> = array_try_from_fn(|i| i.try_into());
26/// assert_eq!(array, Ok([0, 1, 2, 3, 4]));
27///
28/// let array: Result<[i8; 200], _> = array_try_from_fn(|i| i.try_into());
29/// assert!(array.is_err());
30/// ```
31//
32// this is a reimplementation of array::try_from_fn, replace once that became stable
33pub fn array_try_from_fn<E, T, const N: usize>(
34    mut cb: impl FnMut(usize) -> Result<T, E>,
35) -> Result<[T; N], E> {
36    let mut result: MaybeUninit<[T; N]> = MaybeUninit::uninit();
37    {
38        struct DropGuard<'a, T> {
39            slice: &'a mut [MaybeUninit<T>],
40            initialized: usize,
41        }
42        impl<T> Drop for DropGuard<'_, T> {
43            fn drop(&mut self) {
44                for slot in self.slice[..self.initialized].iter_mut() {
45                    // SAFETY: self.initialized is the number of valid elements at all time
46                    // we can assume init and drop them directly
47                    unsafe {
48                        slot.assume_init_drop();
49                    }
50                }
51            }
52        }
53        let mut guard = DropGuard {
54            slice: result.as_mut(),
55            initialized: 0,
56        };
57        for (i, slot) in guard.slice.iter_mut().enumerate() {
58            slot.write(cb(i)?);
59            guard.initialized = i + 1;
60        }
61        // don't drop valid elements
62        guard.initialized = 0;
63    }
64    // SAFETY: All N elements have been successfully written to here
65    unsafe { Ok(result.assume_init()) }
66}
67
68#[cfg(test)]
69mod test {
70    use super::array_try_from_fn;
71    use core::cell::Cell;
72    use std_alloc::rc::Rc;
73    use std_alloc::string::{String, ToString};
74
75    // original test from the documentation
76    #[test]
77    fn array_try_from_fn_test() {
78        let array: Result<[u8; 5], _> = array_try_from_fn(|i| i.try_into());
79        assert_eq!(array, Ok([0, 1, 2, 3, 4]));
80
81        let array: Result<[i8; 200], _> = array_try_from_fn(|i| i.try_into());
82        assert!(array.is_err());
83    }
84
85    #[test]
86    fn smoke_try_from_fn() {
87        let arr = array_try_from_fn(|i| Ok::<_, ()>(i * 2)).unwrap();
88        assert_eq!(arr, [0, 2, 4, 6, 8]);
89        assert_eq!(
90            array_try_from_fn::<_, _, 3>(|i| if i == 0 { Ok(0) } else { Err(1) }).unwrap_err(),
91            1
92        )
93    }
94
95    #[test]
96    fn try_from_fn_dont_drop_on_success() {
97        let arr = array_try_from_fn(|i| Ok::<_, String>(i.to_string())).unwrap();
98        assert_eq!(arr, ["0", "1"]);
99    }
100
101    #[test]
102    fn try_from_fn_drop_on_failure() {
103        let drops = Rc::new(Cell::new(0));
104
105        struct DropCounter(Rc<Cell<usize>>);
106        impl Drop for DropCounter {
107            fn drop(&mut self) {
108                self.0.set(self.0.get() + 1);
109            }
110        }
111
112        let err = array_try_from_fn::<_, _, 10>(|i| match i {
113            0..=4 => Ok(DropCounter(drops.clone())),
114            _ => Err("error".to_string()),
115        })
116        .err()
117        .unwrap();
118        assert_eq!(err, "error");
119        assert_eq!(drops.get(), 5);
120    }
121
122    #[test]
123    #[cfg(feature = "std")]
124    fn try_from_fn_drop_on_panic() {
125        let drops = Rc::new(Cell::new(0));
126
127        struct DropCounter(Rc<Cell<usize>>);
128        impl Drop for DropCounter {
129            fn drop(&mut self) {
130                self.0.set(self.0.get() + 1);
131            }
132        }
133
134        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
135            array_try_from_fn::<_, _, 10>(|i| match i {
136                0..=4 => Ok::<_, String>(DropCounter(drops.clone())),
137                _ => panic!("hi"),
138            })
139        }));
140        assert!(result.is_err());
141        assert_eq!(drops.get(), 5);
142    }
143}
wasmtime_internal_core/array.rs

wasmtime_internal_core/
array.rs
