yggdrasil/kernel/driver/virtio/gpu/src/command.rs

use core::mem::MaybeUninit;

use bytemuck::{Pod, Zeroable};
use device_api::dma::DmaAllocator;
use libk::{
    device::display::PixelFormat,
    dma::{BusAddress, DmaBuffer},
    error::Error,
};
use libk_util::sync::IrqSafeSpinlockGuard;
use ygg_driver_virtio_core::{queue::VirtQueue, transport::Transport};

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ControlHeader {
    pub ty: u32,
    pub flags: u32,
    pub fence_id: u64,
    pub ctx_id: u32,
    pub ring_idx: u8,
    _0: [u8; 3],
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct Rect {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ScanoutInfo {
    pub r: Rect,
    pub enabled: u32,
    pub flags: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ResourceCreate2d {
    pub header: ControlHeader,
    pub resource_id: u32,
    pub format: u32,
    pub width: u32,
    pub height: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct SetScanout {
    pub header: ControlHeader,
    pub r: Rect,
    pub scanout_id: u32,
    pub resource_id: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct GpuMemEntry {
    pub addr: u64,
    pub length: u32,
    pub _0: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ResourceAttachBacking {
    pub header: ControlHeader,
    pub resource_id: u32,
    pub nr_entries: u32,
    pub entry: GpuMemEntry,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ResourceFlush {
    pub header: ControlHeader,
    pub r: Rect,
    pub resource_id: u32,
    pub _0: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct TransferToHost2d {
    pub header: ControlHeader,
    pub r: Rect,
    pub offset: u64,
    pub resource_id: u32,
    pub _0: u32,
}

pub struct ControlLock<'a, T: Transport> {
    control: IrqSafeSpinlockGuard<'a, VirtQueue>,
    transport: IrqSafeSpinlockGuard<'a, T>,
}

impl<'a, T: Transport> ControlLock<'a, T> {
    pub const fn new(
        control: IrqSafeSpinlockGuard<'a, VirtQueue>,
        transport: IrqSafeSpinlockGuard<'a, T>,
    ) -> Self {
        Self { control, transport }
    }

    fn send_recv<'r, Req: Pod>(
        &mut self,
        dma: &dyn DmaAllocator,
        req: &Req,
        buffer: &'r mut DmaBuffer<[MaybeUninit<u8>]>,
    ) -> Result<(&'r ControlHeader, &'r [u8]), Error> {
        let request = DmaBuffer::new_uninit_slice(dma, size_of::<Req>())?;
        let mut request = unsafe { DmaBuffer::assume_init_slice(request) };
        request.copy_from_slice(bytemuck::bytes_of(req));

        let len = self
            .control
            .add_notify_wait_pop(&[buffer], &[&request], &mut *self.transport)
            .inspect_err(|error| {
                log::warn!("virtio queue: {error:?}");
            })
            .map_err(|_| Error::InvalidArgument)? as usize;

        if len < size_of::<ControlHeader>() {
            log::warn!("virtio-gpu: invalid device response length: {len}");
            return Err(Error::InvalidArgument);
        }

        let payload = unsafe { MaybeUninit::slice_assume_init_ref(&buffer[..len]) };
        let header = bytemuck::from_bytes(&payload[..size_of::<ControlHeader>()]);
        let data = &payload[size_of::<ControlHeader>()..len];

        Ok((header, data))
    }

    fn send_recv_no_data<Req: Pod>(
        &mut self,
        dma: &dyn DmaAllocator,
        req: &Req,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
    ) -> Result<(), Error> {
        let (response, _) = self.send_recv(dma, req, buffer)?;
        if response.ty == 0x1100 {
            Ok(())
        } else {
            Err(Error::InvalidArgument)
        }
    }

    pub fn query_scanouts<'r>(
        &mut self,
        dma: &dyn DmaAllocator,
        max_scanouts: usize,
        buffer: &'r mut DmaBuffer<[MaybeUninit<u8>]>,
    ) -> Result<&'r [ScanoutInfo], Error> {
        let request = ControlHeader {
            ty: 0x0100,
            flags: 0,
            fence_id: 0,
            ctx_id: 0,
            ring_idx: 0,
            _0: [0; 3],
        };

        let (response, data) = self.send_recv(dma, &request, buffer)?;
        if response.ty != 0x1101 {
            log::warn!("virtio-gpu: query_scanouts returned {:#x}", response.ty);
            return Err(Error::InvalidArgument);
        }
        let limit = size_of::<ScanoutInfo>() * max_scanouts;

        let scanouts =
            bytemuck::try_cast_slice(&data[..limit]).map_err(|_| Error::InvalidArgument)?;

        Ok(scanouts)
    }

    pub fn create_resource_2d(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        width: u32,
        height: u32,
        pixel_format: PixelFormat,
    ) -> Result<u32, Error> {
        let format = match pixel_format {
            PixelFormat::R8G8B8A8 => 67,
            _ => todo!(),
        };
        let request = ResourceCreate2d {
            header: ControlHeader {
                ty: 0x0101,
                ..ControlHeader::zeroed()
            },
            resource_id: 1,
            width,
            height,
            format,
        };

        self.send_recv_no_data(dma, &request, buffer)?;

        Ok(1)
    }

    pub fn attach_backing(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        resource_id: u32,
        base: BusAddress,
        length: u32,
    ) -> Result<(), Error> {
        let request = ResourceAttachBacking {
            header: ControlHeader {
                ty: 0x0106,
                ..ControlHeader::zeroed()
            },
            resource_id,
            nr_entries: 1,
            entry: GpuMemEntry {
                addr: base.into_u64(),
                length,
                _0: 0,
            },
        };

        self.send_recv_no_data(dma, &request, buffer)
    }

    pub fn set_scanout(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        scanout_id: u32,
        resource_id: u32,
        width: u32,
        height: u32,
    ) -> Result<(), Error> {
        let request = SetScanout {
            header: ControlHeader {
                ty: 0x0103,
                ..ControlHeader::zeroed()
            },
            r: Rect {
                x: 0,
                y: 0,
                width,
                height,
            },
            scanout_id,
            resource_id,
        };

        self.send_recv_no_data(dma, &request, buffer)
    }

    pub fn transfer_to_host_2d(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        resource_id: u32,
        r: Rect,
    ) -> Result<(), Error> {
        let request = TransferToHost2d {
            header: ControlHeader {
                ty: 0x0105,
                ..ControlHeader::zeroed()
            },
            r,
            offset: 0,
            resource_id,
            _0: 0,
        };

        self.send_recv_no_data(dma, &request, buffer)
    }

    pub fn resource_flush(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        resource_id: u32,
        r: Rect,
    ) -> Result<(), Error> {
        let request = ResourceFlush {
            header: ControlHeader {
                ty: 0x0104,
                ..ControlHeader::zeroed()
            },
            r,
            resource_id,
            _0: 0,
        };

        self.send_recv_no_data(dma, &request, buffer)
    }
}