yggdrasil/kernel/driver/usb/xhci/src/ring/transfer.rs

use core::{
    mem::{size_of, MaybeUninit},
    sync::atomic::{AtomicBool, AtomicU64, Ordering},
};

use alloc::{collections::BTreeMap, sync::Arc, vec::Vec};
use bytemuck::{Pod, Zeroable};
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress},
    PageBox,
};
use libk_util::sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock, IrqSafeSpinlockGuard};
use ygg_driver_usb::{
    communication::UsbInterruptTransfer, error::UsbError, pipe::control::ControlTransferSetup,
    UsbControlTransfer, UsbDirection, UsbTransferStatus, UsbTransferToken,
};
use yggdrasil_abi::define_bitfields;

use crate::ring::LinkTrb;

use super::{CommandExecutor, GenericRing, GenericTransferRing};

struct TransferRingInner {
    trbs: PageBox<[MaybeUninit<RawTransferTrb>]>,
    enqueue_index: usize,
    dequeue_index: usize,
    cycle_bit: bool,
}

// TODO split TransferRing into Normal, Control, etc
pub struct ControlTransferRing {
    inner: IrqSafeSpinlock<TransferRingInner>,
    capacity: usize,

    // TODO this is inefficient and ugly
    pending_trbs: IrqSafeRwLock<BTreeMap<PhysicalAddress, UsbTransferToken>>,
    completions: IrqSafeRwLock<BTreeMap<UsbTransferToken, Arc<UsbTransferStatus>>>,

    slot_id: u8,
    ep_id: u8,

    transfer_id: AtomicU64,
    shutdown: AtomicBool,
}

pub struct InterruptInTransferRing {
    inner: IrqSafeSpinlock<TransferRingInner>,
    capacity: usize,

    completions: IrqSafeRwLock<BTreeMap<PhysicalAddress, Arc<UsbTransferStatus>>>,

    slot_id: u8,
    ep_id: u8,

    shutdown: AtomicBool,
}

struct TransferBuilder<'a> {
    ring: &'a ControlTransferRing,
    ring_inner: IrqSafeSpinlockGuard<'a, TransferRingInner>,

    token: UsbTransferToken,
    direction: UsbDirection,
    addresses: Vec<PhysicalAddress>,
    status: Arc<UsbTransferStatus>,
}

impl<'a> TransferBuilder<'a> {
    pub fn new(ring: &'a ControlTransferRing, direction: UsbDirection) -> Self {
        let ring_inner = ring.inner.lock();
        let token = UsbTransferToken(ring.transfer_id.fetch_add(1, Ordering::AcqRel));
        let status = Arc::new(UsbTransferStatus::new());

        ring.completions.write().insert(token, status.clone());

        Self {
            ring,
            ring_inner,

            token,
            direction,
            status,
            addresses: Vec::new(),
        }
    }

    pub fn push_trb<C: TransferTrb>(&mut self, trb: C) -> &mut Self {
        let address = self.ring_inner.enqueue(trb);
        self.addresses.push(address);
        self.ring.pending_trbs.write().insert(address, self.token);
        self
    }

    pub fn start<E: CommandExecutor>(self, executor: &E, length: usize) -> UsbControlTransfer {
        executor.ring_doorbell(self.ring.slot_id as _, self.ring.ep_id);

        UsbControlTransfer {
            id: self.token,
            length,

            direction: self.direction,
            elements: self.addresses,
            status: self.status,
        }
    }
}

impl TransferRingInner {
    fn enqueue<C: TransferTrb>(&mut self, trb: C) -> PhysicalAddress {
        if (self.enqueue_index + 1) % (self.trbs.len() - 1) == self.dequeue_index {
            todo!("Ring full");
        }

        let mut raw: RawTransferTrb = bytemuck::cast(trb);

        raw.flags.set_ty(C::TRB_TYPE as _);
        raw.flags.set_cycle(self.cycle_bit);

        self.trbs[self.enqueue_index].write(raw);

        let address = unsafe { self.trbs.as_physical_address() }
            .add(self.enqueue_index * size_of::<RawTransferTrb>());

        // Move to the next TRB slot
        self.enqueue_index += 1;
        if self.enqueue_index >= self.trbs.len() - 1 {
            self.enqueue_link();

            // Wrap around
            self.cycle_bit = !self.cycle_bit;
            self.enqueue_index = 0;
        }

        address
    }

    fn enqueue_link(&mut self) {
        let base = unsafe { self.trbs.as_physical_address() };

        let link = LinkTrb::new(base, self.cycle_bit);
        self.trbs[self.enqueue_index].write(bytemuck::cast(link));
    }

    fn advance(&mut self) {
        self.dequeue_index += 1;

        if self.dequeue_index >= self.trbs.len() - 1 {
            self.dequeue_index = 0;
        }
    }
}

impl GenericRing for ControlTransferRing {
    fn base(&self) -> PhysicalAddress {
        unsafe { self.inner.lock().trbs.as_physical_address() }
    }

    fn capacity(&self) -> usize {
        self.capacity
    }
}

impl GenericTransferRing for ControlTransferRing {
    fn dequeue_pointer(&self) -> PhysicalAddress {
        let inner = self.inner.lock();
        unsafe { inner.trbs.as_physical_address() }
            .add(inner.dequeue_index * size_of::<RawTransferTrb>())
    }

    fn notify(&self, address: PhysicalAddress, value: u32) {
        if value == 0 {
            return;
        }

        let completions = self.completions.read();
        if let Some(&token) = self.pending_trbs.read().get(&address) {
            let Some(status) = completions.get(&token) else {
                log::warn!(
                    "Notification received for non-existent transfer: {:?}",
                    token
                );
                return;
            };

            status.signal(value);
        }
    }

    fn shutdown(&self) {
        self.shutdown.store(true, Ordering::Release);
        let mut completions = self.completions.write();
        while let Some((_, status)) = completions.pop_first() {
            status.abort();
        }
    }

    fn slot_id(&self) -> u8 {
        self.slot_id
    }

    fn endpoint_id(&self) -> u8 {
        self.ep_id
    }
}

impl GenericRing for InterruptInTransferRing {
    fn base(&self) -> PhysicalAddress {
        unsafe { self.inner.lock().trbs.as_physical_address() }
    }

    fn capacity(&self) -> usize {
        self.capacity
    }
}

impl GenericTransferRing for InterruptInTransferRing {
    fn dequeue_pointer(&self) -> PhysicalAddress {
        let inner = self.inner.lock();
        unsafe { inner.trbs.as_physical_address() }
            .add(inner.dequeue_index * size_of::<RawTransferTrb>())
    }

    fn notify(&self, address: PhysicalAddress, value: u32) {
        if value == 0 {
            return;
        }

        let mut completions = self.completions.write();
        if let Some(status) = completions.remove(&address) {
            status.signal(value);
        }
    }

    fn shutdown(&self) {
        self.shutdown.store(true, Ordering::Release);
        let mut completions = self.completions.write();
        while let Some((_, status)) = completions.pop_first() {
            status.abort();
        }
    }

    fn slot_id(&self) -> u8 {
        self.slot_id
    }

    fn endpoint_id(&self) -> u8 {
        self.ep_id
    }
}

impl InterruptInTransferRing {
    pub fn new(slot_id: u8, ep_id: u8, capacity: usize) -> Result<Self, UsbError> {
        let trbs = PageBox::new_zeroed_slice(capacity).map_err(UsbError::MemoryError)?;

        Ok(Self {
            inner: IrqSafeSpinlock::new(TransferRingInner {
                trbs,
                enqueue_index: 0,
                dequeue_index: 0,
                cycle_bit: true,
            }),
            completions: IrqSafeRwLock::new(BTreeMap::new()),
            slot_id,
            ep_id,
            capacity,
            shutdown: AtomicBool::new(false),
        })
    }

    pub fn start_transfer<E: CommandExecutor>(
        &self,
        executor: &E,
        buffer: &mut PageBox<[u8]>,
    ) -> Result<UsbInterruptTransfer, UsbError> {
        // Don't even try to start the transfer
        if self.shutdown.load(Ordering::Acquire) {
            return Err(UsbError::DeviceDisconnected);
        }

        let status = Arc::new(UsbTransferStatus::new());
        let address = self.inner.lock().enqueue(NormalTransferTrb::new(
            unsafe { buffer.as_physical_address() },
            buffer.len(),
            true,
        ));
        self.completions.write().insert(address, status.clone());

        executor.ring_doorbell(self.slot_id as _, self.ep_id);

        Ok(UsbInterruptTransfer {
            length: buffer.len(),
            direction: UsbDirection::In,
            address,
            status,
        })
    }

    pub fn complete_transfer(&self, _transfer: UsbInterruptTransfer) {
        // Interrupt transfers consist of one TRB each
        // TODO: Can two transfers happen simultaneously? e.g.
        //
        // [TRBa, TRBb] are queued in the ring, both are executing and
        // TRBb finishes first
        self.inner.lock().advance();
    }
}

impl ControlTransferRing {
    pub fn new(slot_id: u8, ep_id: u8, capacity: usize) -> Result<Self, UsbError> {
        let trbs = PageBox::new_zeroed_slice(capacity).map_err(UsbError::MemoryError)?;

        Ok(Self {
            inner: IrqSafeSpinlock::new(TransferRingInner {
                trbs,
                enqueue_index: 0,
                dequeue_index: 0,
                cycle_bit: true,
            }),
            completions: IrqSafeRwLock::new(BTreeMap::new()),
            pending_trbs: IrqSafeRwLock::new(BTreeMap::new()),
            slot_id,
            ep_id,
            capacity,

            transfer_id: AtomicU64::new(0),

            shutdown: AtomicBool::new(false),
        })
    }

    pub fn start_transfer<E: CommandExecutor>(
        &self,
        executor: &E,
        setup: ControlTransferSetup,
        buffer: Option<(PhysicalAddress, usize, UsbDirection)>,
    ) -> Result<UsbControlTransfer, UsbError> {
        // Don't even try to start the transfer
        if self.shutdown.load(Ordering::Acquire) {
            return Err(UsbError::DeviceDisconnected);
        }

        let mut builder = TransferBuilder::new(self, UsbDirection::In);

        builder.push_trb(ControlTransferSetupTrb::new(setup));
        if let Some((address, length, direction)) = buffer {
            builder.push_trb(ControlTransferDataTrb::new(address, length, direction));
        }
        builder.push_trb(ControlTransferStatusTrb::new(UsbDirection::In, true));

        let transfer = builder.start(executor, 0);

        Ok(transfer)
    }

    pub fn complete_transfer(&self, transfer: UsbControlTransfer) {
        let mut pending = self.pending_trbs.write();
        let mut inner = self.inner.lock();
        for trb in transfer.elements {
            pending.remove(&trb);
            inner.advance();
        }
        self.completions.write().remove(&transfer.id);
    }
}

// TRB implementations

define_bitfields! {
    pub RawTransferFlags : u32 {
        (10..16) => ty + set_ty,
        0 => cycle + set_cycle
    }
}

define_bitfields! {
    pub NormalTransferFlags: u64 {
        (0..16) => trb_length,
        37 => interrupt_on_completion,
    }
}

define_bitfields! {
    pub ControlTransferSetupRequest : u64 {
        (0..8) => bm_request_type,
        (8..16) => b_request,
        (16..32) => w_value,
        (32..48) => w_index,
        (48..64) => w_length
    }
}

define_bitfields! {
    pub ControlTransferSetupFlags : u64 {
        (0..16) => trb_length,
        38 => immediate_data,
        (48..50) => transfer_type
    }
}

define_bitfields! {
    pub ControlTransferDataFlags : u64 {
        (0..16) => trb_length,
        48 => direction,
    }
}

define_bitfields! {
    pub ControlTransferStatusFlags : u32 {
        16 => direction,
        5 => interrupt_on_completion
    }
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct NormalTransferTrb {
    pub buffer: PhysicalAddress,
    pub flags: NormalTransferFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct ControlTransferSetupTrb {
    pub request: ControlTransferSetupRequest,
    pub flags: ControlTransferSetupFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct ControlTransferDataTrb {
    pub buffer: PhysicalAddress,
    pub flags: ControlTransferDataFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct ControlTransferStatusTrb {
    _0: [u32; 3],
    pub flags: ControlTransferStatusFlags,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, align(16))]
pub struct RawTransferTrb {
    _0: [u32; 3],
    pub flags: RawTransferFlags,
}

pub trait TransferTrb: Pod {
    const TRB_TYPE: u8;
}

impl NormalTransferTrb {
    pub fn new(buffer: PhysicalAddress, length: usize, interrupt_on_completion: bool) -> Self {
        Self {
            buffer,
            flags: NormalTransferFlags::new(length.try_into().unwrap(), interrupt_on_completion),
        }
    }
}

impl ControlTransferSetupTrb {
    pub const fn new(setup: ControlTransferSetup) -> Self {
        Self {
            request: ControlTransferSetupRequest::new(
                setup.bm_request_type as _,
                setup.b_request as _,
                setup.w_value as _,
                setup.w_index as _,
                setup.w_length as _,
            ),
            flags: ControlTransferSetupFlags::new(8, true, 3),
        }
    }
}

impl ControlTransferDataTrb {
    pub fn new(buffer: PhysicalAddress, length: usize, direction: UsbDirection) -> Self {
        Self {
            buffer,
            flags: ControlTransferDataFlags::new(
                length.try_into().unwrap(),
                direction.is_device_to_host(),
            ),
        }
    }
}

impl ControlTransferStatusTrb {
    pub const fn new(direction: UsbDirection, interrupt_on_completion: bool) -> Self {
        Self {
            _0: [0; 3],
            flags: ControlTransferStatusFlags::new(
                direction.is_device_to_host(),
                interrupt_on_completion,
            ),
        }
    }
}

impl TransferTrb for NormalTransferTrb {
    const TRB_TYPE: u8 = 1;
}

impl TransferTrb for ControlTransferSetupTrb {
    const TRB_TYPE: u8 = 2;
}

impl TransferTrb for ControlTransferDataTrb {
    const TRB_TYPE: u8 = 3;
}

impl TransferTrb for ControlTransferStatusTrb {
    const TRB_TYPE: u8 = 4;
}