axplat-dyn 0.6.2

use alloc::{
    string::{String, ToString},
    vec::Vec,
};
#[cfg(feature = "irq")]
use alloc::{
    sync::{Arc, Weak},
    task::Wake,
};
#[cfg(feature = "irq")]
use core::{
    future::Future,
    pin::pin,
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll, Waker},
};

use ax_driver_base::{BaseDriverOps, DevError, DevResult, DeviceType};
use ax_driver_block::BlockDriverOps;
#[cfg(feature = "irq")]
use ax_kspin::SpinNoIrq;
#[cfg(feature = "irq")]
use ax_plat::irq;
use rd_block::BlkError;
use rdrive::Device;
use spin::Mutex;

use super::DmaImpl;

#[cfg(feature = "phytium-mci")]
mod phytium_mci;
#[cfg(any(feature = "rockchip-sdhci", feature = "rockchip-dwmmc"))]
mod rockchip;
mod virtio;
mod virtio_pci;

pub struct Block {
    name: String,
    irq_num: Option<usize>,
    #[cfg(feature = "irq")]
    irq_state: Option<Arc<BlockIrqState>>,
    queue: Mutex<rd_block::CmdQueue>,
}

pub struct PlatformBlockDevice {
    name: String,
    block: Option<rd_block::Block>,
    irq_num: Option<usize>,
}

impl PlatformBlockDevice {
    fn new(name: String, block: rd_block::Block, irq_num: Option<usize>) -> Self {
        Self {
            name,
            block: Some(block),
            irq_num,
        }
    }
}

impl rdrive::DriverGeneric for PlatformBlockDevice {
    fn name(&self) -> &str {
        &self.name
    }
}

#[cfg(feature = "irq")]
struct BlockIrqState {
    handler: rd_block::IrqHandler,
}

#[cfg(feature = "irq")]
impl BlockIrqState {
    fn handle_irq(&self) {
        self.handler.handle();
    }
}

#[cfg(feature = "irq")]
struct BlockIrqWaiter {
    woken: AtomicBool,
}

#[cfg(feature = "irq")]
impl Wake for BlockIrqWaiter {
    fn wake(self: Arc<Self>) {
        self.wake_by_ref();
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.woken.store(true, Ordering::Release);
    }
}

#[cfg(feature = "irq")]
const BLOCK_IRQ_SLOTS: usize = 16;
#[cfg(feature = "irq")]
static IRQ_SOURCES: [SpinNoIrq<Option<Weak<BlockIrqState>>>; BLOCK_IRQ_SLOTS] =
    [const { SpinNoIrq::new(None) }; BLOCK_IRQ_SLOTS];
#[cfg(feature = "irq")]
const BLOCK_IRQ_REPOLL_SPINS: usize = 256;

impl Block {
    fn use_irq_completion(&self) -> bool {
        #[cfg(feature = "irq")]
        {
            self.irq_state.is_some()
        }

        #[cfg(not(feature = "irq"))]
        {
            false
        }
    }

    fn read_blocks_wait(
        queue: &mut rd_block::CmdQueue,
        block_id: usize,
        block_count: usize,
        use_irq: bool,
    ) -> Vec<Result<rd_block::BlockData, BlkError>> {
        #[cfg(feature = "irq")]
        {
            if use_irq {
                return wait_on_block_irq(queue.read_blocks(block_id, block_count));
            }
        }
        #[cfg(not(feature = "irq"))]
        let _ = use_irq;

        queue.read_blocks_blocking(block_id, block_count)
    }

    fn write_blocks_wait(
        queue: &mut rd_block::CmdQueue,
        block_id: usize,
        data: &[u8],
        use_irq: bool,
    ) -> Vec<Result<(), BlkError>> {
        #[cfg(feature = "irq")]
        {
            if use_irq {
                return wait_on_block_irq(queue.write_blocks(block_id, data));
            }
        }
        #[cfg(not(feature = "irq"))]
        let _ = use_irq;

        queue.write_blocks_blocking(block_id, data)
    }
}

impl BaseDriverOps for Block {
    fn device_type(&self) -> DeviceType {
        DeviceType::Block
    }

    fn device_name(&self) -> &str {
        &self.name
    }

    fn irq_num(&self) -> Option<usize> {
        self.irq_num
    }
}

impl BlockDriverOps for Block {
    fn num_blocks(&self) -> u64 {
        self.queue.lock().num_blocks() as _
    }

    fn block_size(&self) -> usize {
        self.queue.lock().block_size()
    }

    fn flush(&mut self) -> DevResult {
        Ok(())
    }

    fn read_block(&mut self, block_id: u64, buf: &mut [u8]) -> DevResult {
        let block_size = self.block_size();
        if !buf.len().is_multiple_of(block_size) {
            return Err(DevError::InvalidParam);
        }

        let use_irq = self.use_irq_completion();
        let mut queue = self.queue.lock();
        for (offset, chunk) in buf.chunks_mut(block_size).enumerate() {
            let mut blocks =
                Self::read_blocks_wait(&mut queue, block_id as usize + offset, 1, use_irq);
            let block = blocks
                .pop()
                .ok_or(DevError::Io)?
                .map_err(maping_blk_err_to_dev_err)?;
            if block.len() != chunk.len() {
                return Err(DevError::Io);
            }
            chunk.copy_from_slice(&block);
        }
        Ok(())
    }

    fn write_block(&mut self, block_id: u64, buf: &[u8]) -> DevResult {
        let block_size = self.block_size();
        if !buf.len().is_multiple_of(block_size) {
            return Err(DevError::InvalidParam);
        }

        let use_irq = self.use_irq_completion();
        let mut queue = self.queue.lock();
        for (offset, chunk) in buf.chunks(block_size).enumerate() {
            let blocks =
                Self::write_blocks_wait(&mut queue, block_id as usize + offset, chunk, use_irq);
            for block in blocks {
                block.map_err(maping_blk_err_to_dev_err)?;
            }
        }
        Ok(())
    }
}

impl TryFrom<Device<PlatformBlockDevice>> for Block {
    type Error = DevError;

    fn try_from(base: Device<PlatformBlockDevice>) -> Result<Self, Self::Error> {
        let mut dev = base.lock().map_err(|_| DevError::BadState)?;
        let name = dev.name.clone();
        let irq_num = dev.irq_num;
        let mut block = dev.block.take().ok_or(DevError::BadState)?;
        #[cfg(feature = "irq")]
        let irq_handler = irq_num.map(|_| block.irq_handler());
        let queue = block.create_queue().ok_or(DevError::BadState)?;
        drop(dev);

        #[cfg(feature = "irq")]
        let irq_state = if let (Some(irq_num), Some(handler)) = (irq_num, irq_handler) {
            let state = Arc::new(BlockIrqState { handler });
            if let Some(slot) = reserve_block_irq_slot(&state) {
                if irq::register(irq_num, BLOCK_IRQ_HANDLERS[slot]) {
                    Some(state)
                } else {
                    release_block_irq_slot(slot, &state);
                    warn!("failed to register block irq handler for irq {irq_num}");
                    None
                }
            } else {
                warn!("no free block irq source slot for irq {irq_num}");
                None
            }
        } else {
            None
        };
        #[cfg(feature = "irq")]
        let irq_num = irq_state.as_ref().and(irq_num);
        #[cfg(not(feature = "irq"))]
        let irq_num = {
            let _ = irq_num;
            None
        };

        Ok(Self {
            name,
            irq_num,
            #[cfg(feature = "irq")]
            irq_state,
            queue: Mutex::new(queue),
        })
    }
}

pub trait PlatformDeviceBlock {
    fn register_block<T: rd_block::Interface>(self, dev: T);
    fn register_block_with_irq<T: rd_block::Interface>(self, dev: T, irq_num: Option<usize>);
}

impl PlatformDeviceBlock for rdrive::PlatformDevice {
    fn register_block<T: rd_block::Interface>(self, dev: T) {
        self.register_block_with_irq(dev, None);
    }

    fn register_block_with_irq<T: rd_block::Interface>(self, dev: T, irq_num: Option<usize>) {
        #[cfg(feature = "irq")]
        let mut dev = dev;
        #[cfg(feature = "irq")]
        if irq_num.is_some() {
            dev.enable_irq();
        }
        #[cfg(not(feature = "irq"))]
        let dev = {
            let _ = irq_num;
            dev
        };
        let name = dev.name().to_string();
        let dev = rd_block::Block::new(dev, &DmaImpl);
        #[cfg(feature = "irq")]
        let registered_irq_num = irq_num;
        #[cfg(not(feature = "irq"))]
        let registered_irq_num = None;
        self.register(PlatformBlockDevice::new(name, dev, registered_irq_num));
    }
}

#[cfg(any(
    feature = "rockchip-sdhci",
    feature = "rockchip-dwmmc",
    feature = "phytium-mci"
))]
pub(super) fn decode_fdt_irq(interrupts: &[rdrive::probe::fdt::InterruptRef]) -> Option<usize> {
    let interrupt = interrupts.first()?;
    decode_irq_cells(&interrupt.specifier)
}

#[cfg(any(
    feature = "rockchip-sdhci",
    feature = "rockchip-dwmmc",
    feature = "phytium-mci"
))]
fn decode_irq_cells(specifier: &[u32]) -> Option<usize> {
    match specifier {
        [irq] => Some(*irq as usize),
        [kind, irq, ..] => match *kind {
            0 => Some(*irq as usize + 32),
            1 => Some(*irq as usize + 16),
            _ => Some(*irq as usize),
        },
        _ => None,
    }
}

#[cfg(feature = "irq")]
fn handle_block_irq(slot: usize) {
    let Some(state) = IRQ_SOURCES[slot].lock().as_ref().and_then(Weak::upgrade) else {
        return;
    };
    state.handle_irq();
}

#[cfg(feature = "irq")]
fn handle_block_irq_slot<const SLOT: usize>(_: usize) {
    handle_block_irq(SLOT);
}

#[cfg(feature = "irq")]
const BLOCK_IRQ_HANDLERS: [fn(usize); BLOCK_IRQ_SLOTS] = [
    handle_block_irq_slot::<0>,
    handle_block_irq_slot::<1>,
    handle_block_irq_slot::<2>,
    handle_block_irq_slot::<3>,
    handle_block_irq_slot::<4>,
    handle_block_irq_slot::<5>,
    handle_block_irq_slot::<6>,
    handle_block_irq_slot::<7>,
    handle_block_irq_slot::<8>,
    handle_block_irq_slot::<9>,
    handle_block_irq_slot::<10>,
    handle_block_irq_slot::<11>,
    handle_block_irq_slot::<12>,
    handle_block_irq_slot::<13>,
    handle_block_irq_slot::<14>,
    handle_block_irq_slot::<15>,
];

#[cfg(feature = "irq")]
fn reserve_block_irq_slot(state: &Arc<BlockIrqState>) -> Option<usize> {
    for (slot, source) in IRQ_SOURCES.iter().enumerate() {
        let mut source = source.lock();
        if source.as_ref().and_then(Weak::upgrade).is_none() {
            *source = Some(Arc::downgrade(state));
            return Some(slot);
        }
    }
    None
}

#[cfg(feature = "irq")]
fn release_block_irq_slot(slot: usize, state: &Arc<BlockIrqState>) {
    let mut source = IRQ_SOURCES[slot].lock();
    if source
        .as_ref()
        .and_then(Weak::upgrade)
        .is_some_and(|registered| Arc::ptr_eq(&registered, state))
    {
        *source = None;
    }
}

#[cfg(feature = "irq")]
fn wait_on_block_irq<F: Future>(future: F) -> F::Output {
    let mut future = pin!(future);
    let waiter = Arc::new(BlockIrqWaiter {
        woken: AtomicBool::new(false),
    });
    let waker = Waker::from(waiter.clone());
    let mut cx = Context::from_waker(&waker);

    loop {
        waiter.woken.store(false, Ordering::Release);
        match future.as_mut().poll(&mut cx) {
            Poll::Ready(output) => return output,
            Poll::Pending => {
                // rd-block registers the waker after a Retry result, so keep a
                // bounded repoll fallback for completions racing that window.
                for _ in 0..BLOCK_IRQ_REPOLL_SPINS {
                    if waiter.woken.swap(false, Ordering::AcqRel) {
                        break;
                    }
                    core::hint::spin_loop();
                }
            }
        }
    }
}

fn maping_blk_err_to_dev_err(err: BlkError) -> DevError {
    match err {
        BlkError::NotSupported => DevError::Unsupported,
        BlkError::Retry => DevError::Again,
        BlkError::NoMemory => DevError::NoMemory,
        BlkError::InvalidBlockIndex(_) => DevError::InvalidParam,
        BlkError::Other(error) => {
            error!("Block device error: {error}");
            DevError::Io
        }
    }
}