mircuda 0.1.0

use std::marker::PhantomData;

use crate::{DeviceBuffer, DeviceElement, Error, Result, Stream};

/// CUDA grid, block, and dynamic shared-memory geometry.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct LaunchConfig {
    /// Grid dimensions in blocks.
    pub grid: (u32, u32, u32),
    /// Block dimensions in threads.
    pub block: (u32, u32, u32),
    /// Dynamic shared memory requested per block.
    pub shared_memory_bytes: u32,
}

impl LaunchConfig {
    /// Creates a one-dimensional launch covering `elements`.
    pub fn for_elements(elements: usize, threads: u32) -> Result<Self> {
        if elements == 0 || threads == 0 {
            return Err(Error::InvalidLaunch);
        }
        let Ok(elements) = u32::try_from(elements) else {
            return Err(Error::InvalidLaunch);
        };
        Ok(Self {
            grid: (elements.div_ceil(threads), 1, 1),
            block: (threads, 1, 1),
            shared_memory_bytes: 0,
        })
    }

    pub(crate) const fn native(self) -> mircuda_sys::LaunchConfig {
        mircuda_sys::LaunchConfig {
            grid: self.grid,
            block: self.block,
            shared_memory_bytes: self.shared_memory_bytes,
        }
    }
}

/// Unsafe declaration of a CUDA symbol's complete argument contract.
///
/// # Safety
///
/// `NAME` and `Arguments` must exactly match the compiled CUDA symbol. Use
/// [`crate::cuda_export`] instead of implementing this trait manually.
pub unsafe trait KernelSignature: Send + Sync + 'static {
    /// Exported CUDA symbol.
    const NAME: &'static str;
    /// Rust representation of one complete launch argument list.
    type Arguments<'a>;

    /// Encodes typed Rust arguments for the private CUDA launch boundary.
    #[doc(hidden)]
    fn encode(arguments: Self::Arguments<'_>) -> KernelArguments<'_>;
}

/// A loaded CUDA kernel with a statically declared signature.
pub struct TypedKernel<S: KernelSignature> {
    native: mircuda_sys::Kernel,
    signature: PhantomData<S>,
}

impl<S: KernelSignature> Clone for TypedKernel<S> {
    fn clone(&self) -> Self {
        Self {
            native: self.native.clone(),
            signature: PhantomData,
        }
    }
}

impl<S: KernelSignature> std::fmt::Debug for TypedKernel<S> {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("TypedKernel")
            .field("symbol", &S::NAME)
            .finish_non_exhaustive()
    }
}

impl<S: KernelSignature> TypedKernel<S> {
    pub(crate) const fn new(native: mircuda_sys::Kernel) -> Self {
        Self { native, signature: PhantomData }
    }

    pub(crate) const fn native(&self) -> &mircuda_sys::Kernel {
        &self.native
    }

    /// Enqueues a typed launch without synchronizing the host.
    pub fn launch(
        &self,
        stream: &Stream,
        config: LaunchConfig,
        arguments: S::Arguments<'_>,
    ) -> Result<()> {
        let mut arguments = S::encode(arguments);
        Ok(self.native.launch(&stream.native, config.native(), &mut arguments.native)?)
    }

    /// Enqueues a typed launch and returns its exact node from an active capture.
    pub fn launch_captured(
        &self,
        stream: &Stream,
        config: LaunchConfig,
        arguments: S::Arguments<'_>,
    ) -> Result<crate::KernelNode<S>> {
        self.launch(stream, config, arguments)?;
        Ok(crate::KernelNode::new(stream.native.captured_kernel_node(&self.native)?))
    }
}

/// Internal encoded argument list generated by [`crate::cuda_export`].
#[doc(hidden)]
pub struct KernelArguments<'a> {
    pub(crate) native: Vec<mircuda_sys::KernelArgument>,
    borrow: PhantomData<&'a mut ()>,
}

impl<'a> KernelArguments<'a> {
    /// Creates an empty generated argument list.
    #[doc(hidden)]
    #[must_use]
    pub const fn new() -> Self {
        Self { native: Vec::new(), borrow: PhantomData }
    }

    /// Adds an immutable device buffer argument.
    #[doc(hidden)]
    #[must_use]
    pub fn read<T: DeviceElement>(mut self, buffer: &'a DeviceBuffer<T>) -> Self {
        self.native.push(buffer.native.argument());
        self
    }

    /// Adds a mutable device buffer argument.
    #[doc(hidden)]
    #[must_use]
    pub fn write<T: DeviceElement>(mut self, buffer: &'a mut DeviceBuffer<T>) -> Self {
        self.native.push(buffer.native.argument());
        self
    }

    /// Adds a by-value CUDA scalar argument.
    #[doc(hidden)]
    #[must_use]
    pub fn scalar<T: KernelScalar>(mut self, value: T) -> Self {
        self.native.push(value.encode().native());
        self
    }
}

impl Default for KernelArguments<'_> {
    fn default() -> Self {
        Self::new()
    }
}

mod scalar_sealed {
    pub trait Sealed {}
}

/// Plain scalar accepted by generated CUDA signatures.
pub trait KernelScalar: scalar_sealed::Sealed + Copy + Send + Sync + 'static {
    /// Encodes this scalar for the private launch boundary.
    #[doc(hidden)]
    fn encode(self) -> ScalarValue;
}

/// Internal scalar representation used by generated signatures.
#[doc(hidden)]
pub enum ScalarValue {
    U8(u8),
    I8(i8),
    U16(u16),
    I16(i16),
    U32(u32),
    I32(i32),
    U64(u64),
    I64(i64),
    Usize(usize),
    Isize(isize),
    F32(f32),
    F64(f64),
}

impl ScalarValue {
    const fn native(self) -> mircuda_sys::KernelArgument {
        match self {
            Self::U8(value) => mircuda_sys::KernelArgument::U8(value),
            Self::I8(value) => mircuda_sys::KernelArgument::I8(value),
            Self::U16(value) => mircuda_sys::KernelArgument::U16(value),
            Self::I16(value) => mircuda_sys::KernelArgument::I16(value),
            Self::U32(value) => mircuda_sys::KernelArgument::U32(value),
            Self::I32(value) => mircuda_sys::KernelArgument::I32(value),
            Self::U64(value) => mircuda_sys::KernelArgument::U64(value),
            Self::I64(value) => mircuda_sys::KernelArgument::I64(value),
            Self::Usize(value) => mircuda_sys::KernelArgument::Usize(value),
            Self::Isize(value) => mircuda_sys::KernelArgument::Isize(value),
            Self::F32(value) => mircuda_sys::KernelArgument::F32(value),
            Self::F64(value) => mircuda_sys::KernelArgument::F64(value),
        }
    }
}

macro_rules! kernel_scalars {
    ($($scalar:ty => $variant:ident),+ $(,)?) => {$ (
        impl scalar_sealed::Sealed for $scalar {}
        impl KernelScalar for $scalar {
            fn encode(self) -> ScalarValue {
                ScalarValue::$variant(self)
            }
        }
    )+ };
}

kernel_scalars!(
    u8 => U8, i8 => I8, u16 => U16, i16 => I16, u32 => U32, i32 => I32,
    u64 => U64, i64 => I64, usize => Usize, isize => Isize, f32 => F32, f64 => F64,
);