competitive/tools/
array.rs

1#[macro_export]
2macro_rules! array {
3    [@inner $data:ident = [$init:expr; $len:expr]] => {{
4        use ::std::mem::{ManuallyDrop, MaybeUninit};
5        let mut $data: [MaybeUninit<_>; $len] = unsafe { MaybeUninit::uninit().assume_init() };
6        $init;
7        #[repr(C)]
8        union __Transmuter<const N: usize, T: Clone> {
9            src: ManuallyDrop<[MaybeUninit<T>; N]>,
10            dst: ManuallyDrop<[T; N]>,
11        }
12        ManuallyDrop::into_inner(unsafe { __Transmuter { src: ManuallyDrop::new($data) }.dst })
13    }};
14    [|| $e:expr; $len:expr] => {
15        $crate::array![@inner data = [data.iter_mut().for_each(|item| *item = MaybeUninit::new($e)); $len]]
16    };
17    [|$i:pat_param| $e:expr; $len:expr] => {
18        $crate::array![@inner data = [data.iter_mut().enumerate().for_each(|($i, item)| *item = MaybeUninit::new($e)); $len]]
19    };
20    [$e:expr; $len:expr] => {{
21        let e = $e;
22        $crate::array![|| Clone::clone(&e); $len]
23    }};
24}
25
26#[test]
27fn test_array() {
28    let mut x = 0;
29    assert_eq!(array![1; 3], [1; 3]);
30    assert_eq!(array![|| { x += 1; x }; 3], [1, 2, 3]);
31    assert_eq!(array![|i| i + 1; 3], [1, 2, 3]);
32}