90 lines · rust
1// SPDX-License-Identifier: GPL-2.02 3//! Extensions to [`Box`] for fallible allocations.4 5use super::{AllocError, Flags};6use alloc::boxed::Box;7use core::{mem::MaybeUninit, ptr, result::Result};8 9/// Extensions to [`Box`].10pub 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 /// Drops the contents, but keeps the allocation.22 ///23 /// # Examples24 ///25 /// ```26 /// use kernel::alloc::{flags, box_ext::BoxExt};27 /// let value = Box::new([0; 32], flags::GFP_KERNEL)?;28 /// assert_eq!(*value, [0; 32]);29 /// let mut value = Box::drop_contents(value);30 /// // Now we can re-use `value`:31 /// value.write([1; 32]);32 /// // SAFETY: We just wrote to it.33 /// let value = unsafe { value.assume_init() };34 /// assert_eq!(*value, [1; 32]);35 /// # Ok::<(), Error>(())36 /// ```37 fn drop_contents(this: Self) -> Box<MaybeUninit<T>>;38}39 40impl<T> BoxExt<T> for Box<T> {41 fn new(x: T, flags: Flags) -> Result<Self, AllocError> {42 let mut b = <Self as BoxExt<_>>::new_uninit(flags)?;43 b.write(x);44 // SAFETY: We just wrote to it.45 Ok(unsafe { b.assume_init() })46 }47 48 #[cfg(any(test, testlib))]49 fn new_uninit(_flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {50 Ok(Box::new_uninit())51 }52 53 #[cfg(not(any(test, testlib)))]54 fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {55 let ptr = if core::mem::size_of::<MaybeUninit<T>>() == 0 {56 core::ptr::NonNull::<_>::dangling().as_ptr()57 } else {58 let layout = core::alloc::Layout::new::<MaybeUninit<T>>();59 60 // SAFETY: Memory is being allocated (first arg is null). The only other source of61 // safety issues is sleeping on atomic context, which is addressed by klint. Lastly,62 // the type is not a SZT (checked above).63 let ptr =64 unsafe { super::allocator::krealloc_aligned(core::ptr::null_mut(), layout, flags) };65 if ptr.is_null() {66 return Err(AllocError);67 }68 69 ptr.cast::<MaybeUninit<T>>()70 };71 72 // SAFETY: For non-zero-sized types, we allocate above using the global allocator. For73 // zero-sized types, we use `NonNull::dangling`.74 Ok(unsafe { Box::from_raw(ptr) })75 }76 77 fn drop_contents(this: Self) -> Box<MaybeUninit<T>> {78 let ptr = Box::into_raw(this);79 // SAFETY: `ptr` is valid, because it came from `Box::into_raw`.80 unsafe { ptr::drop_in_place(ptr) };81 82 // CAST: `MaybeUninit<T>` is a transparent wrapper of `T`.83 let ptr = ptr.cast::<MaybeUninit<T>>();84 85 // SAFETY: `ptr` is valid for writes, because it came from `Box::into_raw` and it is valid for86 // reads, since the pointer came from `Box::into_raw` and the type is `MaybeUninit<T>`.87 unsafe { Box::from_raw(ptr) }88 }89}90