yggdrasil/kernel/driver/block/ahci/src/port.rs

use core::{
    future::poll_fn,
    mem::MaybeUninit,
    sync::atomic::{AtomicU32, Ordering},
    task::{Context, Poll},
};

use alloc::{boxed::Box, string::String};
use async_trait::async_trait;
use bytemuck::Zeroable;
use futures_util::task::AtomicWaker;
use libk::vfs::block::NgBlockDevice;
use libk_mm::{address::AsPhysicalAddress, device::DeviceMemoryIo, PageBox, PageSlice};
use libk_util::{sync::IrqSafeSpinlock, waker::QueueWaker, OneTimeInit};
use tock_registers::interfaces::{Readable, Writeable};

use crate::{
    command::{AtaCommand, AtaIdentify, AtaReadDmaEx},
    data::{CommandListEntry, CommandTable, ReceivedFis, COMMAND_LIST_LENGTH},
    error::AhciError,
    regs::{PortRegs, CMD_PENDING, CMD_READY, IE, TFD},
    AhciController, MAX_COMMANDS, SECTOR_SIZE,
};

#[derive(Clone, Copy, PartialEq, Debug)]
pub enum PortType {
    Sata,
}

struct PortInner {
    regs: DeviceMemoryIo<'static, PortRegs>,

    #[allow(unused)]
    received_fis: PageBox<ReceivedFis>,
    command_list: PageBox<[CommandListEntry]>,
}

pub struct PortInfo {
    pub model: String,
    pub serial: String,
    pub lba_count: u64,
}

#[allow(unused)]
pub struct AhciPort {
    inner: IrqSafeSpinlock<PortInner>,
    ahci: &'static AhciController,
    ty: PortType,
    pub(crate) index: usize,
    info: OneTimeInit<PortInfo>,

    command_allocation: IrqSafeSpinlock<u32>,
    // One command index can only be waited for by one task, so this approach is usable
    command_completion: [(AtomicWaker, AtomicU32); COMMAND_LIST_LENGTH],
    command_available: QueueWaker,
}

struct SubmittedCommand<'a> {
    port: &'a AhciPort,
    index: usize,
}

impl<'a> SubmittedCommand<'a> {
    pub async fn wait_for_completion(self) -> Result<(), AhciError> {
        let result = poll_fn(|cx| self.port.poll_slot(cx, self.index)).await;

        // Free the command without dropping it
        self.port.free_command(self.index);
        core::mem::forget(self);

        result
    }
}

impl<'a> Drop for SubmittedCommand<'a> {
    fn drop(&mut self) {
        panic!(
            "Cannot drop command in flight: port{}, slot{}",
            self.port.index, self.index
        )
    }
}

impl PortInner {
    fn submit_command<C: AtaCommand>(
        &mut self,
        index: usize,
        command: &C,
    ) -> Result<(), AhciError> {
        let list_entry = &mut self.command_list[index];
        let mut table_entry =
            PageBox::new(CommandTable::zeroed()).map_err(AhciError::MemoryError)?;

        table_entry.setup_command(command)?;
        *list_entry = CommandListEntry::new(
            unsafe { table_entry.as_physical_address() },
            command.regions().len(),
        )?;

        // Sync before send
        // XXX do this properly
        #[cfg(target_arch = "x86_64")]
        unsafe {
            core::arch::asm!("wbinvd");
        }

        // TODO deal with this async way
        while self.regs.TFD.matches_any(TFD::BSY::SET + TFD::DRQ::SET) {
            core::hint::spin_loop();
        }

        let ci = self.regs.CI.get();
        assert_eq!(ci & (1 << index), 0);
        self.regs.CI.set(ci | (1 << index));

        Ok(())
    }
}

impl AhciPort {
    pub fn create(
        regs: DeviceMemoryIo<'static, PortRegs>,
        ahci: &'static AhciController,
        index: usize,
    ) -> Result<&'static Self, AhciError> {
        log::debug!("Initialize port {}", index);
        regs.stop()?;

        if !ahci.has_64_bit {
            todo!("Handle controllers incapable of 64 bit");
        }

        let received_fis = PageBox::new(ReceivedFis::zeroed()).map_err(AhciError::MemoryError)?;
        let command_list = PageBox::new_slice(CommandListEntry::zeroed(), COMMAND_LIST_LENGTH)
            .map_err(AhciError::MemoryError)?;

        regs.set_received_fis_address_64(unsafe { received_fis.as_physical_address() });
        regs.set_command_list_address_64(unsafe { command_list.as_physical_address() });

        regs.IE.write(
            IE::DPE::SET
                + IE::IFE::SET
                + IE::OFE::SET
                + IE::HBDE::SET
                + IE::HBFE::SET
                + IE::TFEE::SET
                + IE::DHRE::SET,
        );

        regs.start()?;

        let inner = PortInner {
            regs,
            command_list,
            received_fis,
        };
        let command_completion =
            [const { (AtomicWaker::new(), AtomicU32::new(CMD_READY)) }; MAX_COMMANDS];
        let command_available = QueueWaker::new();
        let command_allocation = IrqSafeSpinlock::new(0);

        Ok(Box::leak(Box::new(Self {
            inner: IrqSafeSpinlock::new(inner),
            ty: PortType::Sata,
            info: OneTimeInit::new(),
            ahci,
            index,

            command_completion,
            command_allocation,
            command_available,
        })))
    }

    pub async fn init(&'static self) -> Result<(), AhciError> {
        let identify = self.perform_command(AtaIdentify::create()?).await?;
        let model = identify.model_number.to_string();
        let serial = identify.serial_number.to_string();
        let lba_count = identify.logical_sector_count();

        // TODO can sector size be different from 512 in ATA?
        //      should logical sector size be accounted for?
        // TODO test for ReadDmaEx capability (?)

        self.info.init(PortInfo {
            model,
            serial,
            lba_count,
        });

        Ok(())
    }

    pub fn info(&self) -> Option<&PortInfo> {
        self.info.try_get()
    }

    async fn allocate_command(&self) -> usize {
        poll_fn(|cx| {
            self.command_available.register(cx.waker());

            let mut state = self.command_allocation.lock();
            if *state != u32::MAX {
                self.command_available.remove(cx.waker());

                for i in 0..MAX_COMMANDS {
                    if *state & (1 << i) == 0 {
                        *state |= 1 << i;
                        self.command_completion[i]
                            .1
                            .store(CMD_PENDING, Ordering::Release);
                        return Poll::Ready(i);
                    }
                }

                unreachable!()
            } else {
                Poll::Pending
            }
        })
        .await
    }

    async fn submit<C: AtaCommand>(&self, command: &C) -> Result<SubmittedCommand, AhciError> {
        let index = self.allocate_command().await;
        if let Err(error) = self.inner.lock().submit_command(index, command) {
            self.free_command(index);
            return Err(error);
        }
        Ok(SubmittedCommand { port: self, index })
    }

    async fn perform_command<C: AtaCommand>(&self, command: C) -> Result<C::Response, AhciError> {
        // Run the command
        self.submit(&command).await?.wait_for_completion().await?;

        Ok(unsafe { command.into_response() })
    }

    fn poll_slot(&self, cx: &mut Context<'_>, index: usize) -> Poll<Result<(), AhciError>> {
        let (waker, status) = &self.command_completion[index];

        match status.load(Ordering::Acquire) {
            CMD_PENDING => (),
            CMD_READY => return Poll::Ready(Ok(())),
            _ => return Poll::Ready(Err(AhciError::DeviceError)),
        }

        waker.register(cx.waker());

        match status.load(Ordering::Acquire) {
            CMD_PENDING => Poll::Pending,
            CMD_READY => Poll::Ready(Ok(())),
            _ => Poll::Ready(Err(AhciError::DeviceError)),
        }
    }

    fn free_command(&self, index: usize) {
        {
            let mut alloc = self.command_allocation.lock();
            assert_ne!(*alloc & (1 << index), 0);
            *alloc &= !(1 << index);
        }
        self.command_available.wake_one();
    }

    pub fn handle_pending_interrupts(&self) -> bool {
        let inner = self.inner.lock();

        let Some(status) = inner.regs.clear_interrupt() else {
            return false;
        };

        let ci = inner.regs.CI.get();

        for i in 0..MAX_COMMANDS {
            if ci & (1 << i) == 0 {
                if self.command_completion[i]
                    .1
                    .swap(status.into(), Ordering::Release)
                    == CMD_PENDING
                {
                    log::info!("port{}: completion on slot {}", self.index, i);
                    self.command_completion[i].0.wake();
                }
            }
        }

        true
    }
}

#[async_trait]
impl NgBlockDevice for AhciPort {
    type Error = AhciError;

    async fn read(
        &self,
        lba: u64,
        buffer: &mut PageSlice<MaybeUninit<u8>>,
    ) -> Result<(), AhciError> {
        let command = AtaReadDmaEx::new(lba, buffer.len() / SECTOR_SIZE, buffer);
        self.submit(&command).await?.wait_for_completion().await
    }

    async fn write(&self, _lba: u64, _buffer: &PageSlice<u8>) -> Result<(), AhciError> {
        // TODO AtaDmaWriteEx
        Err(AhciError::FeatureNotImplemented)
    }

    fn block_size(&self) -> usize {
        SECTOR_SIZE
    }

    fn block_count(&self) -> usize {
        self.info().as_ref().map(|i| i.lba_count).unwrap() as _
    }

    fn max_blocks_per_request(&self) -> usize {
        // TODO
        1
    }
}