rust traits for converting iterators to arrays with no extra allocations

to_array.rs

#![feature(min_const_generics,maybe_uninit_extra)]

// VERY UNTESTED AND NOT REVIEWED
// THIS CODE USES UNSAFE RUST AND COULD LEAD TO UB

use std::mem::{self, MaybeUninit};

#[derive(Clone,Debug)]
enum ToArrayError {
    TooShort(usize, usize),
    TooLong(usize)
}

trait ToArray<T> {
    /// Convert to an array with zero allocations
    fn take_array<const N: usize>(&mut self) -> Result<[T; N], ToArrayError>;
    fn to_array<const N: usize>(self) -> Result<[T; N], ToArrayError>;
}

impl<I, T: Sized> ToArray<T> for I where I: Iterator<Item=T> {
    fn take_array<const N: usize>(&mut self) -> Result<[T; N], ToArrayError> {
        let mut res: [MaybeUninit<T>; N] = unsafe {
            MaybeUninit::uninit().assume_init()
        };
        
        let mut error_index = None;
        
        for (i, el) in res.iter_mut().enumerate() {
            if let Some(x) = self.next() {
                *el = MaybeUninit::new(x);
            } else {
                error_index = Some(i);
                break;
            }
        }
        
        if let Some(i) = error_index {
            // drop initialized elements
            for el in &mut res[..i] {
                unsafe { el.assume_init_drop() };
            } 
            Err(ToArrayError::TooShort(i, N))
        } else {
            Ok(unsafe {
                mem::transmute_copy(&res)
            })
        }
    } 
    
    fn to_array<const N: usize>(mut self) -> Result<[T; N], ToArrayError> {
        let arr = self.take_array()?;
        match self.next() {
            Some(_) => Err(ToArrayError::TooLong(N)),
            None => Ok(arr)
        } 
    }
}

trait ToArrayDefault<T> {
    fn take_array_default<const N: usize>(&mut self) -> [T; N];
    fn to_array_default<const N: usize>(self) -> Result<[T; N], ToArrayError>;
}

impl<I, T: Sized + Default> ToArrayDefault<T> for I where I: Iterator<Item=T> {
    fn take_array_default<const N: usize>(&mut self) -> [T; N] {
        let mut res: [MaybeUninit<T>; N] = unsafe {
            MaybeUninit::uninit().assume_init()
        };
        
        for el in &mut res {
            *el = MaybeUninit::new(self.next().unwrap_or_else(|| Default::default()));
        }
        unsafe {
            mem::transmute_copy(&res)
        }
    } 
    
    fn to_array_default<const N: usize>(mut self) -> Result<[T; N], ToArrayError> {
        let arr = self.take_array_default();
        match self.next() {
            Some(_) => Err(ToArrayError::TooLong(N)),
            None => Ok(arr)
        } 
    }
}

trait ToArrayPad<T> {
    fn take_array_pad<const N: usize>(&mut self, pad: T) -> [T; N];
    fn to_array_pad<const N: usize>(self, pad: T) -> Result<[T; N], ToArrayError>;
}

impl<I, T: Sized + Clone> ToArrayPad<T> for I where I: Iterator<Item=T> {
    fn take_array_pad<const N: usize>(&mut self, pad: T) -> [T; N] {
        let mut res: [MaybeUninit<T>; N] = unsafe {
            MaybeUninit::uninit().assume_init()
        };
        
        for el in &mut res {
            *el = MaybeUninit::new(self.next().unwrap_or_else(|| pad.clone()));
        }
        unsafe {
            mem::transmute_copy(&res)
        }
    } 
    
    fn to_array_pad<const N: usize>(mut self, pad: T) -> Result<[T; N], ToArrayError> {
        let arr = self.take_array_pad(pad);
        match self.next() {
            Some(_) => Err(ToArrayError::TooLong(N)),
            None => Ok(arr)
        } 
    }
}


fn main() {
    let arr1: [i32; 60] = (0..60).to_array().unwrap();
    let arr2: [i32; 60] = (0..50).to_array_default().unwrap();
    println!("{:?}", arr1);
    println!("{:?}", arr2);
}