arbitrary/foreign/core/
array.rs

1use {
2    crate::{size_hint, Arbitrary, Result, Unstructured},
3    core::{
4        array,
5        mem::{self, MaybeUninit},
6        ptr,
7    },
8};
9
10/// Helper to safely create arrays since the standard library doesn't
11/// provide one yet. Shouldn't be necessary in the future.
12struct ArrayGuard<T, const N: usize> {
13    dst: *mut T,
14    initialized: usize,
15}
16
17impl<T, const N: usize> Drop for ArrayGuard<T, N> {
18    fn drop(&mut self) {
19        debug_assert!(self.initialized <= N);
20        let initialized_part = ptr::slice_from_raw_parts_mut(self.dst, self.initialized);
21        unsafe {
22            ptr::drop_in_place(initialized_part);
23        }
24    }
25}
26
27fn try_create_array<F, T, const N: usize>(mut cb: F) -> Result<[T; N]>
28where
29    F: FnMut(usize) -> Result<T>,
30{
31    let mut array: MaybeUninit<[T; N]> = MaybeUninit::uninit();
32    let array_ptr = array.as_mut_ptr();
33    let dst = array_ptr as _;
34    let mut guard: ArrayGuard<T, N> = ArrayGuard {
35        dst,
36        initialized: 0,
37    };
38    unsafe {
39        for (idx, value_ptr) in (*array.as_mut_ptr()).iter_mut().enumerate() {
40            ptr::write(value_ptr, cb(idx)?);
41            guard.initialized += 1;
42        }
43        mem::forget(guard);
44        Ok(array.assume_init())
45    }
46}
47
48impl<'a, T, const N: usize> Arbitrary<'a> for [T; N]
49where
50    T: Arbitrary<'a>,
51{
52    #[inline]
53    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
54        try_create_array(|_| <T as Arbitrary<'a>>::arbitrary(u))
55    }
56
57    #[inline]
58    fn arbitrary_take_rest(mut u: Unstructured<'a>) -> Result<Self> {
59        let mut array = Self::arbitrary(&mut u)?;
60        if let Some(last) = array.last_mut() {
61            *last = Arbitrary::arbitrary_take_rest(u)?;
62        }
63        Ok(array)
64    }
65
66    #[inline]
67    fn size_hint(depth: usize) -> (usize, Option<usize>) {
68        Self::try_size_hint(depth).unwrap_or_default()
69    }
70
71    #[inline]
72    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
73        let hint = <T as Arbitrary>::try_size_hint(depth)?;
74        Ok(size_hint::and_all(&array::from_fn::<_, N, _>(|_| hint)))
75    }
76}