yggdrasil/lib/kernel-util/src/mem/mod.rs

use core::{
    alloc::Layout,
    fmt,
    mem::{size_of, MaybeUninit},
    ops::{Deref, DerefMut},
};

use yggdrasil_abi::error::Error;

use crate::api::{__allocate_contiguous_pages, __free_page, __physicalize};

use self::address::{AsPhysicalAddress, PhysicalAddress};

pub mod address;
pub mod device;
pub mod pointer;
pub mod table;

pub struct PageBox<T: ?Sized> {
    value: *mut T,
    page_count: usize,
}

impl<T> PageBox<T> {
    #[inline]
    fn alloc_slice(count: usize) -> Result<(PhysicalAddress, usize), Error> {
        // TODO hardcoded page sizes
        let layout = Layout::array::<T>(count).unwrap();
        let page_count = (layout.size() + 0xFFF) / 0x1000;
        Ok((
            unsafe { __allocate_contiguous_pages(page_count) }?,
            page_count,
        ))
    }

    #[inline]
    fn alloc() -> Result<(PhysicalAddress, usize), Error> {
        let page_count = (size_of::<T>() + 0xFFF) / 0x1000;
        Ok((
            unsafe { __allocate_contiguous_pages(page_count) }?,
            page_count,
        ))
    }

    pub fn new(init: T) -> Result<PageBox<T>, Error> {
        let (base, page_count) = Self::alloc()?;
        let value = base.virtualize_raw() as *mut T;

        unsafe {
            value.write(init);
        }

        let result = PageBox { value, page_count };
        result.trace_created();
        Ok(result)
    }

    pub fn new_slice(item: T, count: usize) -> Result<PageBox<[T]>, Error>
    where
        T: Copy,
    {
        let (base, page_count) = Self::alloc_slice(count)?;
        let base_virt_ptr = base.virtualize_raw() as *mut T;
        let value = core::ptr::slice_from_raw_parts_mut(base_virt_ptr, count);

        for i in 0..count {
            unsafe {
                value.get_unchecked_mut(i).write(item);
            }
        }

        let result = PageBox { value, page_count };
        result.trace_created();
        Ok(result)
    }

    pub fn new_uninit() -> Result<PageBox<MaybeUninit<T>>, Error> {
        let (base, page_count) = PageBox::<MaybeUninit<T>>::alloc()?;
        let value = base.virtualize_raw() as *mut MaybeUninit<T>;
        let result = PageBox { value, page_count };
        result.trace_created();
        Ok(result)
    }

    pub fn new_uninit_slice(count: usize) -> Result<PageBox<[MaybeUninit<T>]>, Error> {
        let (base, page_count) = PageBox::<MaybeUninit<T>>::alloc_slice(count)?;
        let base_virt_ptr = base.virtualize_raw() as *mut MaybeUninit<T>;
        let value = core::ptr::slice_from_raw_parts_mut(base_virt_ptr, count);
        let result = PageBox { value, page_count };
        result.trace_created();
        Ok(result)
    }
}

impl<T: ?Sized> PageBox<T> {
    #[inline]
    pub fn as_ptr(&self) -> *const T {
        self.value as _
    }

    #[inline]
    fn trace_created(&self) {
        log::trace!(
            "Alloc PageBox<{}> @ {:p}, {}",
            core::any::type_name::<T>(),
            self.value,
            self.page_count
        );
    }

    #[inline]
    fn trace_dropped(&self) {
        log::trace!(
            "Free PageBox<{}> @ {:p}, {}",
            core::any::type_name::<T>(),
            self.value,
            self.page_count
        );
    }
}

impl<T> PageBox<MaybeUninit<T>> {
    /// Consumes the [PageBox], returning a new one with [MaybeUninit] removed.
    ///
    /// # Safety
    ///
    /// See [MaybeUninit::assume_init_mut].
    pub unsafe fn assume_init(self) -> PageBox<T> {
        // SAFETY: Memory-safe, as:
        // 1. MaybeUninit<T> is transparent
        // 2. self.value still points to the same memory and is not deallocated
        let page_count = self.page_count;
        let value = MaybeUninit::assume_init_mut(&mut *self.value);

        // Prevent deallocation of the PageBox with MaybeUninit
        core::mem::forget(self);

        PageBox { value, page_count }
    }
}

impl<T> PageBox<[MaybeUninit<T>]> {
    /// Consumes the [PageBox], returning a new one with [MaybeUninit] removed.
    ///
    /// # Safety
    ///
    /// See [MaybeUninit::slice_assume_init_mut].
    pub unsafe fn assume_init_slice(self) -> PageBox<[T]> {
        // SAFETY: Memory-safe, as:
        // 1. MaybeUninit<T> is transparent
        // 2. self.value still points to the same memory and is not deallocated
        let page_count = self.page_count;
        let value = MaybeUninit::slice_assume_init_mut(&mut *self.value);

        core::mem::forget(self);

        PageBox { value, page_count }
    }

    /// Returns a reference to the slice data with [MaybeUninit] removed.
    ///
    /// # Safety
    ///
    /// See [MaybeUninit::slice_assume_init_ref]
    pub unsafe fn assume_init_slice_ref(&self) -> &[T] {
        MaybeUninit::slice_assume_init_ref(self.deref())
    }

    /// Returns a mutable reference to the slice data with [MaybeUninit] removed.
    ///
    /// # Safety
    ///
    /// See [MaybeUninit::slice_assume_init_mut]
    pub unsafe fn assume_init_slice_mut(&mut self) -> &mut [T] {
        MaybeUninit::slice_assume_init_mut(self.deref_mut())
    }
}

impl<T: ?Sized> AsPhysicalAddress for PageBox<T> {
    #[inline]
    unsafe fn as_physical_address(&self) -> PhysicalAddress {
        PhysicalAddress(__physicalize(self.value.addr()))
    }
}

impl<T: ?Sized> Deref for PageBox<T> {
    type Target = T;

    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        unsafe { &*self.value }
    }
}

impl<T: ?Sized> DerefMut for PageBox<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.value }
    }
}

impl<T: ?Sized> Drop for PageBox<T> {
    fn drop(&mut self) {
        self.trace_dropped();
        unsafe {
            core::ptr::drop_in_place(self.value);
        }
        // SAFETY: Safe, pointer obtained through "virtualize"
        let base = PhysicalAddress(unsafe { __physicalize(self.value.addr()) });
        for i in 0..self.page_count {
            // SAFETY: Safe, page allocated only by this PageBox
            unsafe {
                __free_page(base.add(0x1000 * i));
            }
        }
    }
}

impl<T: ?Sized> fmt::Pointer for PageBox<T> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.value.fmt(f)
    }
}

impl<T: ?Sized + fmt::Debug> fmt::Debug for PageBox<T> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self.deref(), f)
    }
}

impl<T: ?Sized + fmt::Display> fmt::Display for PageBox<T> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.deref(), f)
    }
}

unsafe impl<T: ?Sized + Send> Send for PageBox<T> {}
unsafe impl<T: ?Sized + Sync> Sync for PageBox<T> {}