xref: /linux/rust/kernel/alloc/box_ext.rs (revision 502cc061de6692a9a8ca9bcf486de78e2664e869)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 //! Extensions to [`Box`] for fallible allocations.
4 
5 use super::{AllocError, Flags};
6 use alloc::boxed::Box;
7 use core::mem::MaybeUninit;
8 
9 /// Extensions to [`Box`].
10 pub trait BoxExt<T>: Sized {
11     /// Allocates a new box.
12     ///
13     /// The allocation may fail, in which case an error is returned.
14     fn new(x: T, flags: Flags) -> Result<Self, AllocError>;
15 
16     /// Allocates a new uninitialised box.
17     ///
18     /// The allocation may fail, in which case an error is returned.
19     fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError>;
20 }
21 
22 impl<T> BoxExt<T> for Box<T> {
23     fn new(x: T, flags: Flags) -> Result<Self, AllocError> {
24         let mut b = <Self as BoxExt<_>>::new_uninit(flags)?;
25         b.write(x);
26         // SAFETY: We just wrote to it.
27         Ok(unsafe { b.assume_init() })
28     }
29 
30     #[cfg(any(test, testlib))]
31     fn new_uninit(_flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
32         Ok(Box::new_uninit())
33     }
34 
35     #[cfg(not(any(test, testlib)))]
36     fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
37         let ptr = if core::mem::size_of::<MaybeUninit<T>>() == 0 {
38             core::ptr::NonNull::<_>::dangling().as_ptr()
39         } else {
40             let layout = core::alloc::Layout::new::<MaybeUninit<T>>();
41 
42             // SAFETY: Memory is being allocated (first arg is null). The only other source of
43             // safety issues is sleeping on atomic context, which is addressed by klint. Lastly,
44             // the type is not a SZT (checked above).
45             let ptr =
46                 unsafe { super::allocator::krealloc_aligned(core::ptr::null_mut(), layout, flags) };
47             if ptr.is_null() {
48                 return Err(AllocError);
49             }
50 
51             ptr.cast::<MaybeUninit<T>>()
52         };
53 
54         // SAFETY: For non-zero-sized types, we allocate above using the global allocator. For
55         // zero-sized types, we use `NonNull::dangling`.
56         Ok(unsafe { Box::from_raw(ptr) })
57     }
58 }
59