yggdrasil/kernel/libk/src/dma.rs

use core::{
    alloc::Layout,
    fmt,
    mem::{self, MaybeUninit},
    ops::{Deref, DerefMut, Range, Sub},
    ptr::{self, NonNull},
};

use bytemuck::{Pod, Zeroable};
use device_api::dma::{DmaAllocation, DmaAllocator};
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress, Virtualize},
    phys, L3_PAGE_SIZE,
};
use yggdrasil_abi::error::Error;

pub struct DummyDmaAllocator;

impl DmaAllocator for DummyDmaAllocator {
    fn allocate(&self, layout: Layout) -> Result<DmaAllocation, Error> {
        if layout.align() > L3_PAGE_SIZE {
            return Err(Error::InvalidMemoryOperation);
        }
        let page_count = layout.size().div_ceil(L3_PAGE_SIZE);
        let host_physical = phys::alloc_pages_contiguous(page_count)?;
        let host_virtual = unsafe {
            NonNull::new_unchecked(ptr::with_exposed_provenance_mut(host_physical.virtualize()))
        };
        let bus_address = host_physical.into_u64();

        Ok(DmaAllocation {
            host_physical: host_physical.into_u64(),
            host_virtual,
            bus_address,
            page_count,
        })
    }
}

pub struct DmaBuffer<T: ?Sized> {
    host_pointer: NonNull<T>,
    host_physical: PhysicalAddress,
    bus_address: u64,
    page_count: usize,
}

pub struct DmaSlice<'a, T> {
    buffer: &'a DmaBuffer<[T]>,
    range: Range<usize>,
}

pub struct DmaSliceMut<'a, T> {
    buffer: &'a mut DmaBuffer<[T]>,
    range: Range<usize>,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug, PartialOrd, Ord, Pod, Zeroable)]
#[repr(transparent)]
pub struct BusAddress(u64);

impl<T> DmaBuffer<T> {
    pub fn new(allocator: &dyn DmaAllocator, value: T) -> Result<DmaBuffer<T>, Error> {
        let mut uninit = DmaBuffer::new_uninit(allocator)?;
        uninit.write(value);
        Ok(unsafe { DmaBuffer::assume_init(uninit) })
    }

    pub fn new_slice(
        allocator: &dyn DmaAllocator,
        value: T,
        size: usize,
    ) -> Result<DmaBuffer<[T]>, Error>
    where
        T: Copy,
    {
        let mut uninit = DmaBuffer::new_uninit_slice(allocator, size)?;
        for i in 0..size {
            uninit[i].write(value);
        }
        Ok(unsafe { DmaBuffer::assume_init_slice(uninit) })
    }

    pub fn from_slice(allocator: &dyn DmaAllocator, source: &[T]) -> Result<DmaBuffer<[T]>, Error>
    where
        T: Copy,
    {
        let mut uninit = DmaBuffer::new_uninit_slice(allocator, source.len())?;
        uninit[..].write_copy_of_slice(source);
        Ok(unsafe { DmaBuffer::assume_init_slice(uninit) })
    }

    pub fn new_slice_with<F: Fn(usize) -> T>(
        allocator: &dyn DmaAllocator,
        init: F,
        size: usize,
    ) -> Result<DmaBuffer<[T]>, Error> {
        let mut uninit = DmaBuffer::new_uninit_slice(allocator, size)?;
        for i in 0..size {
            uninit[i].write(init(i));
        }
        Ok(unsafe { DmaBuffer::assume_init_slice(uninit) })
    }

    pub fn new_uninit(allocator: &dyn DmaAllocator) -> Result<DmaBuffer<MaybeUninit<T>>, Error> {
        let layout = Layout::new::<T>();
        let allocation = allocator.allocate(layout)?;
        let host_pointer = allocation.host_virtual.cast();
        Ok(DmaBuffer {
            host_pointer,
            host_physical: PhysicalAddress::from_u64(allocation.host_physical),
            bus_address: allocation.bus_address,
            page_count: allocation.page_count,
        })
    }

    pub fn new_uninit_slice(
        allocator: &dyn DmaAllocator,
        size: usize,
    ) -> Result<DmaBuffer<[MaybeUninit<T>]>, Error> {
        let layout = Layout::array::<T>(size).map_err(|_| Error::InvalidMemoryOperation)?;
        let allocation = allocator.allocate(layout)?;
        let host_pointer = NonNull::slice_from_raw_parts(allocation.host_virtual.cast(), size);
        Ok(DmaBuffer {
            host_pointer,
            host_physical: PhysicalAddress::from_u64(allocation.host_physical),
            bus_address: allocation.bus_address,
            page_count: allocation.page_count,
        })
    }

    pub fn new_zeroed_slice(
        allocator: &dyn DmaAllocator,
        size: usize,
    ) -> Result<DmaBuffer<[MaybeUninit<T>]>, Error> {
        let layout = Layout::array::<T>(size).map_err(|_| Error::InvalidMemoryOperation)?;
        let allocation = allocator.allocate(layout)?;
        unsafe {
            let mut slice = NonNull::<[u8]>::slice_from_raw_parts(
                allocation.host_virtual.cast(),
                layout.size(),
            );

            slice.as_mut().fill(0);
        }
        let host_pointer = NonNull::slice_from_raw_parts(allocation.host_virtual.cast(), size);
        Ok(DmaBuffer {
            host_pointer,
            host_physical: PhysicalAddress::from_u64(allocation.host_physical),
            bus_address: allocation.bus_address,
            page_count: allocation.page_count,
        })
    }

    pub unsafe fn assume_init(buffer: DmaBuffer<MaybeUninit<T>>) -> DmaBuffer<T> {
        let host_pointer = buffer.host_pointer;
        let host_physical = buffer.host_physical;
        let bus_address = buffer.bus_address;
        let page_count = buffer.page_count;

        mem::forget(buffer);

        let host_pointer = host_pointer.cast();

        DmaBuffer {
            host_pointer,
            host_physical,
            bus_address,
            page_count,
        }
    }

    pub unsafe fn assume_init_slice(buffer: DmaBuffer<[MaybeUninit<T>]>) -> DmaBuffer<[T]> {
        let host_pointer = buffer.host_pointer;
        let host_physical = buffer.host_physical;
        let bus_address = buffer.bus_address;
        let page_count = buffer.page_count;

        mem::forget(buffer);

        let len = host_pointer.len();
        let host_pointer = NonNull::slice_from_raw_parts(host_pointer.cast::<T>(), len);

        DmaBuffer {
            host_pointer,
            host_physical,
            bus_address,
            page_count,
        }
    }
}

impl<T> DmaBuffer<[T]> {
    fn slice_range_check(&self, range: &Range<usize>) {
        assert!(
            range.end <= self.len() && range.start <= self.len(),
            "DMA buffer slice range out of bounds"
        );
        assert!(range.start <= range.end, "Invalid DMA slice range");
    }

    pub fn slice(&self, range: Range<usize>) -> DmaSlice<'_, T> {
        self.slice_range_check(&range);
        DmaSlice {
            buffer: self,
            range,
        }
    }

    pub fn slice_mut(&mut self, range: Range<usize>) -> DmaSliceMut<'_, T> {
        self.slice_range_check(&range);
        DmaSliceMut {
            buffer: self,
            range,
        }
    }

    pub fn cache_flush_element(&self, index: usize, write: bool) {
        libk_mm::flush_cache_data(&raw const self[index], write);
    }

    pub fn cache_flush_range(&self, range: Range<usize>, write: bool) {
        libk_mm::flush_cache_data_slice(&raw const self[range], write);
    }

    pub fn cache_flush_all(&self, write: bool) {
        libk_mm::flush_cache_data_slice(&raw const self[..], write);
    }
}

impl<T: ?Sized> DmaBuffer<T> {
    #[inline]
    pub fn page_count(&self) -> usize {
        self.page_count
    }

    #[inline]
    pub fn bus_address(&self) -> BusAddress {
        BusAddress(self.bus_address)
    }
}

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

impl<T: ?Sized> Drop for DmaBuffer<T> {
    fn drop(&mut self) {
        log::trace!("Drop DmaBuffer @ {:#x}", self.host_physical);
        unsafe {
            ptr::drop_in_place(self.host_pointer.as_ptr());
            #[cfg(any(rust_analyzer, target_os = "none"))]
            for i in 0..self.page_count {
                phys::free_page(self.host_physical.add(i * L3_PAGE_SIZE));
            }
        }
    }
}

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

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

    fn deref(&self) -> &Self::Target {
        unsafe { self.host_pointer.as_ref() }
    }
}

impl<T: ?Sized> DerefMut for DmaBuffer<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { self.host_pointer.as_mut() }
    }
}

impl<T: ?Sized> fmt::Pointer for DmaBuffer<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "<dma@{:#x}>", self.bus_address)
    }
}

impl BusAddress {
    pub const ZERO: Self = Self(0);

    pub const fn into_u64(self) -> u64 {
        self.0
    }

    pub fn try_into_u32(self) -> Result<u32, Error> {
        self.0.try_into().map_err(|_| Error::InvalidMemoryOperation)
    }

    pub const fn add(self, offset: usize) -> Self {
        Self(self.0 + offset as u64)
    }
}

impl Sub<BusAddress> for BusAddress {
    type Output = usize;

    fn sub(self, rhs: BusAddress) -> Self::Output {
        (self.0 - rhs.0).try_into().unwrap()
    }
}

impl fmt::LowerHex for BusAddress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::LowerHex::fmt(&self.0, f)
    }
}

impl<'a, T> DmaSlice<'a, T> {
    pub fn bus_address(&self) -> BusAddress {
        self.buffer.bus_address().add(self.range.start)
    }

    // TODO subslicing
    pub fn into_parts(self) -> (&'a DmaBuffer<[T]>, Range<usize>) {
        (self.buffer, self.range)
    }

    pub fn cache_flush_all(&self, write: bool) {
        self.buffer.cache_flush_range(self.range.clone(), write);
    }
}

impl<T> Deref for DmaSlice<'_, T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        &self.buffer[self.range.clone()]
    }
}

impl<T> fmt::Pointer for DmaSlice<'_, T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:p}[{:?}]", *self.buffer, self.range)
    }
}

impl<'a, T> DmaSliceMut<'a, T> {
    pub fn bus_address(&self) -> BusAddress {
        self.buffer.bus_address().add(self.range.start)
    }

    // TODO subslicing
    pub fn into_parts(self) -> (&'a mut DmaBuffer<[T]>, Range<usize>) {
        (self.buffer, self.range)
    }

    pub fn cache_flush_all(&self, write: bool) {
        self.buffer.cache_flush_range(self.range.clone(), write);
    }
}

impl<T> Deref for DmaSliceMut<'_, T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        &self.buffer[self.range.clone()]
    }
}

impl<T> DerefMut for DmaSliceMut<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.buffer[self.range.clone()]
    }
}

impl<T> fmt::Pointer for DmaSliceMut<'_, T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:p}[{:?}]", *self.buffer, self.range)
    }
}