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

use std::{
    alloc, mem,
    ptr::{self, NonNull},
};

/// A typesafe helper that stores the allocated pointer without the data initialized.
pub struct BoxAllocation<T>(
    // ptr cannot be null since it would mean the allocation failed.
    // Note: covariance is acceptable since this eventually becomes a `Box<T>`,
    // which is covariant too.
    NonNull<T>,
);

impl<T> BoxAllocation<T> {
    /// Consumes self and writes the given value into the allocation.
    #[inline(always)] // if this does not get inlined then copying happens
    pub fn init(self, value: T) -> Box<T> {
        if mem::size_of::<T>() == 0 {
            return Box::new(value);
        }

        unsafe {
            let ptr = self.0.as_ptr();
            mem::forget(self);
            ptr::write(ptr, value);
            Box::from_raw(ptr)
        }
    }
}

impl<T> Drop for BoxAllocation<T> {
    fn drop(&mut self) {
        if mem::size_of::<T>() == 0 {
            return;
        }

        let layout = alloc::Layout::new::<T>();
        unsafe {
            alloc::dealloc(self.0.as_ptr() as *mut u8, layout);
        }
    }
}

/// Helper trait for a `Box` type that allocates up-front.
pub trait BoxHelper<T> {
    /// Allocates the storage without providing any data.
    fn alloc() -> BoxAllocation<T>;
}

impl<T> BoxHelper<T> for Box<T> {
    fn alloc() -> BoxAllocation<T> {
        if mem::size_of::<T>() == 0 {
            return BoxAllocation(NonNull::dangling());
        }

        let layout = alloc::Layout::new::<T>();
        BoxAllocation(
            NonNull::new(unsafe { alloc::alloc(layout) as *mut T })
                .unwrap_or_else(|| alloc::handle_alloc_error(layout)), // oom
        )
    }
}