• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 use {
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.
12 struct ArrayGuard<T, const N: usize> {
13     dst: *mut T,
14     initialized: usize,
15 }
16 
17 impl<T, const N: usize> Drop for ArrayGuard<T, N> {
drop(&mut self)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 
try_create_array<F, T, const N: usize>(mut cb: F) -> Result<[T; N]> where F: FnMut(usize) -> Result<T>,27 fn try_create_array<F, T, const N: usize>(mut cb: F) -> Result<[T; N]>
28 where
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 
48 impl<'a, T, const N: usize> Arbitrary<'a> for [T; N]
49 where
50     T: Arbitrary<'a>,
51 {
52     #[inline]
arbitrary(u: &mut Unstructured<'a>) -> Result<Self>53     fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
54         try_create_array(|_| <T as Arbitrary<'a>>::arbitrary(u))
55     }
56 
57     #[inline]
arbitrary_take_rest(mut u: Unstructured<'a>) -> Result<Self>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]
size_hint(depth: usize) -> (usize, Option<usize>)67     fn size_hint(depth: usize) -> (usize, Option<usize>) {
68         Self::try_size_hint(depth).unwrap_or_default()
69     }
70 
71     #[inline]
try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached>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 }
77