1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

#![cfg_attr(not(feature = "std"), no_std)]

use alloc::{alloc::alloc, boxed::Box};
use core::{alloc::Layout, mem::size_of, ptr::NonNull};

extern crate alloc;

/// All useful implementations of this trait are already provided by this crate.
///
/// Indicates that the implementing type is an array that can be initialised with a closure.
///
/// The `F` generic type parameter is used to convince Rust to allow this -
/// it is normally (always?) inferred by the compiler.
pub unsafe trait BuildArray<C, F> {
    unsafe fn build_into(space: *mut Self, build_component: &mut impl FnMut(*mut C));
}

unsafe impl<T, const N: usize> BuildArray<T, ()> for [T; N] {
    unsafe fn build_into(space: *mut Self, build_component: &mut impl FnMut(*mut T)) {
        let space = space as *mut T;

        for n in 0..N {
            let container_ptr = space.add(n);
            build_component(container_ptr);
        }
    }
}

unsafe impl<T, Q, F, const N: usize> BuildArray<T, (F,)> for [Q; N]
where
    Q: BuildArray<T, F>,
{
    unsafe fn build_into(space: *mut Self, build_component: &mut impl FnMut(*mut T)) {
        let space = space as *mut Q;

        for n in 0..N {
            let container_ptr = space.add(n);
            Q::build_into(container_ptr, build_component);
        }
    }
}

/// Constructs an array directly on the heap.
///
/// The provided closure is called in memory-order to provide a value for
/// the next slot in the returned array.
///
/// The size and the dimensionality of the returned array is not limited.
pub fn with<T, A, F>(mut construct: impl FnMut() -> T) -> Box<A>
where
    A: BuildArray<T, F>,
{
    if size_of::<A>() == 0 {
        // Safe as A is a ZST.
        return unsafe { Box::from_raw(NonNull::dangling().as_mut()) };
    }

    // Safe as A is not a ZST
    let space = unsafe { alloc(Layout::new::<A>()) };

    #[cfg(all(test, not(miri)))]
    {
        for i in 0..size_of::<A>() {
            // Safe as space is currently just bytes
            unsafe { space.add(i).write(0x69) }
        }
    }

    // Safe as space is allocated according to A's layout
    let space = space as *mut A;

    // Safe as storage is valid and properly aligned
    let mut build_component = move |storage: *mut T| unsafe { storage.write(construct()) };

    unsafe {
        A::build_into(space, &mut build_component);
    }

    // Safe as space is allocated according to Box's rules and is properly initialised
    unsafe { Box::from_raw(space) }
}

/// Constructs an array of default values directly on the heap.
///
/// The size and the dimensionality of the returned array is not limited.
///
/// Due to the fact that [Default] is implemented for (some) arrays
/// where the contained type also provides [Default], the compiler may
/// struggle to infer all the types for this function when a multi-dimensional
/// array is requested. If this is the case, the function can be called
/// as `from_default::<T, _, _>()` where `T` is the type that you wish to
/// use the [Default] implementation of.
pub fn from_default<T, A, F>() -> Box<A>
where
    A: BuildArray<T, F>,
    T: Default,
{
    with(|| Default::default())
}

/// Constructs an array of cloned values directly on the heap.
///
/// The size and the dimensionality of the returned array is not limited.
pub fn from_cloned<T, A, F>(val: &T) -> Box<A>
where
    A: BuildArray<T, F>,
    T: Clone,
{
    with(|| val.clone())
}

#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn test_it_works() {
        let mut n = 5;
        let thing: Box<[[[i32; 5]; 5]; 5]> = with(|| {
            n += 1;
            n
        });

        #[cfg(feature = "std")]
        println!("{:?}", thing);

        let thing2: Box<[[[i32; 5]; 5]; 5]> = from_default::<i32, _, _>();

        #[cfg(feature = "std")]
        println!("{:?}", thing2);
    }
}