vec_utils/
boxed.rs

1use std::alloc::Layout;
2use std::mem::ManuallyDrop;
3use std::ptr::NonNull;
4
5/// Extension methods for `Box<T>`
6pub trait BoxExt: Sized {
7    /// The type that the `Box<T>` stores
8    type T: ?Sized;
9
10    /// drops the value inside the box and returns the allocation
11    /// in the form of an `UninitBox`
12    fn drop_box(bx: Self) -> UninitBox;
13
14    /// takes the value inside the box and returns it as well as the
15    /// allocation in the form of an `UninitBox`
16    fn take_box(bx: Self) -> (UninitBox, Self::T)
17    where
18        Self::T: Sized;
19}
20
21impl<T: ?Sized> BoxExt for Box<T> {
22    type T = T;
23
24    fn drop_box(bx: Self) -> UninitBox {
25        unsafe {
26            let layout = Layout::for_value::<T>(&bx);
27            let ptr = NonNull::new_unchecked(Box::into_raw(bx));
28
29            ptr.as_ptr().drop_in_place();
30
31            UninitBox {
32                ptr: ptr.cast(),
33                layout,
34            }
35        }
36    }
37
38    fn take_box(bx: Self) -> (UninitBox, Self::T)
39    where
40        Self::T: Sized,
41    {
42        unsafe {
43            let ptr = NonNull::new_unchecked(Box::into_raw(bx));
44
45            let value = ptr.as_ptr().read();
46
47            (
48                UninitBox {
49                    ptr: ptr.cast(),
50                    layout: Layout::new::<T>(),
51                },
52                value,
53            )
54        }
55    }
56}
57
58/// An uninitialized piece of memory
59pub struct UninitBox {
60    ptr: NonNull<u8>,
61    layout: Layout,
62}
63
64impl UninitBox {
65    /// The layout of the allocation
66    #[inline]
67    pub fn layout(&self) -> Layout {
68        self.layout
69    }
70
71    /// create a new allocation that can fit the given type
72    #[inline]
73    pub fn new<T>() -> Self {
74        Self::from_layout(Layout::new::<T>())
75    }
76
77    /// Create a new allocation that can fit the given layout
78    #[inline]
79    pub fn from_layout(layout: Layout) -> Self {
80        if layout.size() == 0 {
81            UninitBox {
82                layout,
83                ptr: unsafe { NonNull::new_unchecked(layout.align() as *mut u8) },
84            }
85        } else {
86            let ptr = unsafe { std::alloc::alloc(layout) };
87
88            if ptr.is_null() {
89                std::alloc::handle_alloc_error(layout)
90            } else {
91                unsafe {
92                    UninitBox {
93                        ptr: NonNull::new_unchecked(ptr),
94                        layout,
95                    }
96                }
97            }
98        }
99    }
100
101    /// Initialize the box with the given value,
102    ///
103    /// # Panic
104    ///
105    /// if `std::alloc::Layout::new::<T>() != self.layout()` then
106    /// this function will panic
107    #[inline]
108    pub fn init<T>(self, value: T) -> Box<T> {
109        assert_eq!(
110            self.layout,
111            Layout::new::<T>(),
112            "Layout of UninitBox is incompatible with `T`"
113        );
114
115        let bx = ManuallyDrop::new(self);
116
117        let ptr = bx.ptr.cast::<T>().as_ptr();
118
119        unsafe {
120            ptr.write(value);
121
122            Box::from_raw(ptr)
123        }
124    }
125
126    /// Initialize the box with the given value,
127    ///
128    /// # Panic
129    ///
130    /// if `std::alloc::Layout::new::<T>() != self.layout()` then
131    /// this function will panic
132    #[inline]
133    pub fn init_with<T, F: FnOnce() -> T>(self, value: F) -> Box<T> {
134        assert_eq!(
135            self.layout,
136            Layout::new::<T>(),
137            "Layout of UninitBox is incompatible with `T`"
138        );
139
140        let bx = ManuallyDrop::new(self);
141
142        let ptr = bx.ptr.cast::<T>().as_ptr();
143
144        unsafe {
145            ptr.write(value());
146
147            Box::from_raw(ptr)
148        }
149    }
150
151    /// Get the pointer from the `UninitBox`
152    ///
153    /// This pointer is not valid to write to
154    #[inline]
155    pub fn as_ptr(&self) -> *const () {
156        self.ptr.as_ptr() as *const ()
157    }
158
159    /// Get the pointer from the `UninitBox`
160    #[inline]
161    pub fn as_mut_ptr(&mut self) -> *mut () {
162        self.ptr.as_ptr() as *mut ()
163    }
164}
165
166impl Drop for UninitBox {
167    fn drop(&mut self) {
168        unsafe { std::alloc::dealloc(self.ptr.as_ptr(), self.layout) }
169    }
170}