mlx-c-rust 0.32.4-alpha.2

use crate::MlxCError;
use crate::raw;

/// Owned MLX stream handle used to preserve runtime thread affinity.
#[derive(Debug)]
pub struct MlxStream(raw::mlx_stream);

impl MlxStream {
    pub fn default_cpu() -> Result<Self, MlxCError> {
        // SAFETY: The runtime error handler is installed and the returned
        // stream is placed immediately under RAII ownership.
        let raw_stream = unsafe { raw::mlx_default_cpu_stream_new() };
        Self::from_raw(raw_stream, "acquire the default MLX CPU stream")
    }

    pub fn default_gpu() -> Result<Self, MlxCError> {
        // SAFETY: The runtime error handler is installed and the returned
        // stream is placed immediately under RAII ownership.
        let raw_stream = unsafe { raw::mlx_default_gpu_stream_new() };
        Self::from_raw(raw_stream, "acquire the default MLX GPU stream")
    }

    pub const fn raw(&self) -> raw::mlx_stream {
        self.0
    }

    fn from_raw(raw_stream: raw::mlx_stream, operation: &'static str) -> Result<Self, MlxCError> {
        if raw_stream.ctx.is_null() {
            return Err(MlxCError {
                operation,
                description: "MLX returned an empty stream handle".to_owned(),
            });
        }
        Ok(Self(raw_stream))
    }

    pub(crate) fn from_live_raw(
        raw_stream: raw::mlx_stream,
        operation: &'static str,
    ) -> Result<Self, MlxCError> {
        Self::from_raw(raw_stream, operation)
    }
}

impl Drop for MlxStream {
    fn drop(&mut self) {
        // SAFETY: This owner releases its live stream handle exactly once.
        unsafe {
            raw::mlx_stream_free(self.0);
        }
    }
}