Files
RedBear-OS/src/common/aligned_box.rs
T
2024-02-25 20:06:22 +01:00

123 lines
3.7 KiB
Rust

use core::alloc::{GlobalAlloc, Layout};
use crate::{common::unique::Unique, memory::Enomem};
// Necessary because GlobalAlloc::dealloc requires the layout to be the same, and therefore Box
// cannot be used for increased alignment directly.
// TODO: move to common?
pub struct AlignedBox<T: ?Sized, const ALIGN: usize> {
inner: Unique<T>,
}
pub unsafe trait ValidForZero {}
unsafe impl<const N: usize> ValidForZero for [u8; N] {}
unsafe impl ValidForZero for u8 {}
impl<T: ?Sized, const ALIGN: usize> AlignedBox<T, ALIGN> {
fn layout(&self) -> Layout {
layout_upgrade_align(Layout::for_value::<T>(&*self), ALIGN)
}
}
const fn layout_upgrade_align(layout: Layout, align: usize) -> Layout {
const fn max(a: usize, b: usize) -> usize {
if a > b {
a
} else {
b
}
}
let Ok(x) = Layout::from_size_align(layout.size(), max(align, layout.align())) else {
panic!("failed to calculate layout");
};
x
}
impl<T, const ALIGN: usize> AlignedBox<T, ALIGN> {
#[inline(always)]
pub fn try_zeroed() -> Result<Self, Enomem>
where
T: ValidForZero,
{
Ok(unsafe {
let ptr =
crate::ALLOCATOR.alloc_zeroed(layout_upgrade_align(Layout::new::<T>(), ALIGN));
if ptr.is_null() {
return Err(Enomem);
}
Self {
inner: Unique::new_unchecked(ptr.cast()),
}
})
}
}
impl<T, const ALIGN: usize> AlignedBox<[T], ALIGN> {
#[inline]
pub fn try_zeroed_slice(len: usize) -> Result<Self, Enomem>
where
T: ValidForZero,
{
Ok(unsafe {
let ptr = crate::ALLOCATOR.alloc_zeroed(layout_upgrade_align(
Layout::array::<T>(len).unwrap(),
ALIGN,
));
if ptr.is_null() {
return Err(Enomem);
}
Self {
inner: Unique::new_unchecked(core::ptr::slice_from_raw_parts_mut(ptr.cast(), len)),
}
})
}
}
impl<T: ?Sized, const ALIGN: usize> core::fmt::Debug for AlignedBox<T, ALIGN> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(
f,
"[aligned box at {:p}, size {} alignment {}]",
self.inner.as_ptr(),
self.layout().size(),
self.layout().align()
)
}
}
impl<T: ?Sized, const ALIGN: usize> Drop for AlignedBox<T, ALIGN> {
fn drop(&mut self) {
unsafe {
let layout = self.layout();
core::ptr::drop_in_place(self.inner.as_ptr());
crate::ALLOCATOR.dealloc(self.inner.as_ptr().cast(), layout);
}
}
}
impl<T: ?Sized, const ALIGN: usize> core::ops::Deref for AlignedBox<T, ALIGN> {
type Target = T;
fn deref(&self) -> &Self::Target {
unsafe { &*self.inner.as_ptr() }
}
}
impl<T: ?Sized, const ALIGN: usize> core::ops::DerefMut for AlignedBox<T, ALIGN> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.inner.as_ptr() }
}
}
impl<T: Clone + ValidForZero, const ALIGN: usize> Clone for AlignedBox<T, ALIGN> {
fn clone(&self) -> Self {
let mut new =
Self::try_zeroed().unwrap_or_else(|_| alloc::alloc::handle_alloc_error(self.layout()));
T::clone_from(&mut new, self);
new
}
}
impl<T: Clone + ValidForZero, const ALIGN: usize> Clone for AlignedBox<[T], ALIGN> {
fn clone(&self) -> Self {
let mut new = Self::try_zeroed_slice(self.len())
.unwrap_or_else(|_| alloc::alloc::handle_alloc_error(self.layout()));
for i in 0..self.len() {
new[i].clone_from(&self[i]);
}
new
}
}