mlx-c-rust 0.32.4-alpha.3

use crate::MlxBindingsContext;
use crate::MlxCError;
use crate::{MlxArray, MlxDtype, raw};

impl MlxBindingsContext {
    /// Copies float32 values into a new owned MLX array after validating shape arithmetic.
    pub fn array_from_f32(&self, values: &[f32], shape: &[i32]) -> Result<MlxArray, MlxCError> {
        MlxArray::from_f32(values, shape)
    }

    /// Copies int32 values into a new owned MLX array after validating shape arithmetic.
    pub fn array_from_i32(&self, values: &[i32], shape: &[i32]) -> Result<MlxArray, MlxCError> {
        MlxArray::from_i32(values, shape)
    }

    /// Copies uint32 values into a new owned MLX array after validating shape arithmetic.
    pub fn array_from_u32(&self, values: &[u32], shape: &[i32]) -> Result<MlxArray, MlxCError> {
        MlxArray::from_u32(values, shape)
    }

    /// Creates an int32 half-open range with unit stride.
    pub fn arange_i32(&self, start: i32, stop: i32) -> Result<MlxArray, MlxCError> {
        self.output_array("create an MLX integer range", |output_array, stream| {
            // SAFETY: Scalar arguments are copied and output is uniquely writable.
            unsafe {
                raw::mlx_arange(
                    output_array,
                    f64::from(start),
                    f64::from(stop),
                    1.0,
                    MlxDtype::Int32.to_raw(),
                    stream,
                )
            }
        })
    }

    /// Creates a float32 half-open range with an explicit stride.
    pub fn arange_f32(&self, start: f64, stop: f64, stride: f64) -> Result<MlxArray, MlxCError> {
        self.output_array("create an MLX float32 range", |output_array, stream| {
            // SAFETY: Scalar arguments are copied and output is uniquely writable.
            unsafe {
                raw::mlx_arange(
                    output_array,
                    start,
                    stop,
                    stride,
                    MlxDtype::Float32.to_raw(),
                    stream,
                )
            }
        })
    }

    /// Creates a zero-filled MLX array with a validated static shape.
    pub fn zeros(&self, shape: &[i32], dtype: MlxDtype) -> Result<MlxArray, MlxCError> {
        validate_creation_shape("create a zero-filled MLX array", shape)?;
        self.output_array("create a zero-filled MLX array", |output_array, stream| {
            // SAFETY: Shape remains borrowed for this graph-building call, dtype
            // is a known MLX dtype, stream is live, and output is uniquely writable.
            unsafe {
                raw::mlx_zeros(
                    output_array,
                    shape.as_ptr(),
                    shape.len(),
                    dtype.to_raw(),
                    stream,
                )
            }
        })
    }

    /// Creates a floating MLX array filled by one scalar without host expansion.
    pub fn full(
        &self,
        shape: &[i32],
        fill_value: f32,
        dtype: MlxDtype,
    ) -> Result<MlxArray, MlxCError> {
        const OPERATION: &str = "create a filled MLX array";
        validate_creation_shape(OPERATION, shape)?;
        if !matches!(
            dtype,
            MlxDtype::Float16 | MlxDtype::Float32 | MlxDtype::BFloat16
        ) || !fill_value.is_finite()
        {
            return Err(MlxCError {
                operation: OPERATION,
                description: "fill value must be finite and dtype must be supported floating point"
                    .to_owned(),
            });
        }
        let float_fill_value = self.array_from_f32(&[fill_value], &[])?;
        self.output_array(OPERATION, |output_array, stream| {
            // SAFETY: Shape and scalar remain live for this graph-building call;
            // dtype is known and output is uniquely writable.
            unsafe {
                raw::mlx_full(
                    output_array,
                    shape.as_ptr(),
                    shape.len(),
                    float_fill_value.raw(),
                    dtype.to_raw(),
                    stream,
                )
            }
        })
    }
}

fn validate_creation_shape(operation: &'static str, shape: &[i32]) -> Result<(), MlxCError> {
    if shape.iter().any(|dimension_size| *dimension_size < 0) {
        return Err(MlxCError {
            operation,
            description: "array dimensions must be nonnegative".to_owned(),
        });
    }

    let element_count = shape
        .iter()
        .try_fold(1_usize, |element_count, dimension_size| {
            let dimension_size = usize::try_from(*dimension_size).ok()?;
            element_count.checked_mul(dimension_size)
        });
    if element_count.is_none() {
        return Err(MlxCError {
            operation,
            description: "array element count overflows usize".to_owned(),
        });
    }

    Ok(())
}