Compiler Explorer

Source code

#![feature(maybe_uninit_slice)]
#![feature(maybe_uninit_array_assume_init)]
#![feature(maybe_uninit_extra)]
#![feature(maybe_uninit_uninit_array)]
#![feature(array_map)]

use core::mem::MaybeUninit;
use core::{mem, ptr};

fn collect_into_array<I, const N: usize>(iter: &mut I) -> Option<[I::Item; N]>
where
    I: Iterator,
{
    if N == 0 {
        // TODO: is there a better way to achieve this?

// SAFETY: since N is 0, the array is empty and there are no
        // uninitialized elements.
        let out = unsafe { MaybeUninit::array_assume_init(MaybeUninit::uninit_array::<N>()) };
        return Some(out);
    }

struct Guard<T, const N: usize> {
        ptr: *mut T,
        initialized: usize,
    }

impl<T, const N: usize> Drop for Guard<T, N> {
        fn drop(&mut self) {
            debug_assert!(self.initialized <= N);

let initialized_part =
                crate::ptr::slice_from_raw_parts_mut(self.ptr, self.initialized);

// SAFETY: this raw slice will contain only initialized objects.
            unsafe {
                crate::ptr::drop_in_place(initialized_part);
            }
        }
    }

let mut array = MaybeUninit::uninit_array::<N>();
    let mut guard: Guard<I::Item, N> = Guard {
        ptr: MaybeUninit::slice_as_mut_ptr(&mut array),
        initialized: 0,
    };

while let Some(item) = iter.next() {
        // SAFETY: `guard.initialized` starts at 0, is increased by one in the
        // loop and the loop is aborted once it reaches N (which is
        // `array.len()`). Therefore this index operation is safe.
        unsafe {
            array.get_unchecked_mut(guard.initialized).write(item);
        }
        guard.initialized += 1;

// Check if the whole array was initialized.
        if guard.initialized == N {
            // SAFETY: the condition above asserts that all elements are
            // initialized.
            crate::mem::forget(guard);
            let out = unsafe { MaybeUninit::array_assume_init(array) };
            return Some(out);
        }
    }

// This is only reached if the iterator is exhausted before
    // `guard.initialized` reaches `N`. Also note that `guard` is dropped here,
    // dropping all already initialized elements.
    None
}

fn map<T, F, U, const N: usize>(arr: [T; N], f: F) -> [U; N] 
where
    F: FnMut(T) -> U,
{
    match collect_into_array(&mut std::array::IntoIter::new(arr).map(f)) {
        Some(array) => array,
        None => unsafe { core::hint::unreachable_unchecked() },
    }
}

pub fn map_ints_std(s: [u8; 4]) -> [u32; 4] {
    s.map(|i| i as _)
}

pub fn map_ints_new(s: [u8; 4]) -> [u32; 4] {
    map(s, |i| i as _)
}