mozu-audio 0.1.0

//! Core ML tensor allocation and strided access.

use super::{Error, Result, core_error};
use block2::StackBlock;
#[cfg(test)]
use objc2::rc::autoreleasepool;
use objc2::{AnyThread, rc::Retained};
use objc2_core_ml::{MLMultiArray, MLMultiArrayDataType};
use objc2_foundation::{NSArray, NSNumber};
use std::{cell::RefCell, ffi::c_void, ptr::NonNull};

pub(super) trait Scalar: Copy + Default {
    const TYPE: MLMultiArrayDataType;
}
impl Scalar for f32 {
    const TYPE: MLMultiArrayDataType = MLMultiArrayDataType::Float32;
}
impl Scalar for i32 {
    const TYPE: MLMultiArrayDataType = MLMultiArrayDataType::Int32;
}
impl Scalar for half::f16 {
    const TYPE: MLMultiArrayDataType = MLMultiArrayDataType::Float16;
}

fn shape_product(shape: &[usize]) -> Result<usize> {
    if shape.is_empty() {
        return Err(Error::Invalid("Empty tensor shape".into()));
    }
    shape
        .iter()
        .try_fold(1usize, |n, &d| n.checked_mul(d).filter(|_| d > 0))
        .ok_or_else(|| Error::Invalid(format!("Invalid tensor shape: {shape:?}")))
}

pub(super) fn flat_offset(index: usize, shape: &[usize], strides: &[usize]) -> usize {
    let (mut remainder, mut offset) = (index, 0);
    for i in (0..shape.len()).rev() {
        offset += (remainder % shape[i]) * strides[i];
        remainder /= shape[i];
    }
    offset
}

pub(super) fn array_shape(array: &MLMultiArray) -> Vec<usize> {
    // SAFETY: The retained Core ML array supplies NSNumber dimensions.
    unsafe { array.shape() }
        .iter()
        .map(|n| n.unsignedIntegerValue())
        .collect()
}

struct TensorLayout {
    shape: Vec<usize>,
    strides: Vec<usize>,
    count: usize,
    minimum_bytes: usize,
}

impl TensorLayout {
    fn new<T: Scalar>(array: &MLMultiArray) -> Result<Self> {
        // SAFETY: Metadata belongs to this valid retained array.
        if unsafe { array.dataType() } != T::TYPE {
            return Err(Error::Invalid(format!(
                "Unexpected tensor dtype: expected={:?}, actual={:?}",
                T::TYPE,
                unsafe { array.dataType() }
            )));
        }
        let shape = array_shape(array);
        let count = shape_product(&shape)?;
        // SAFETY: Core ML owns the dimension and stride metadata.
        let strides: Vec<usize> = unsafe { array.strides() }
            .iter()
            .map(|n| n.unsignedIntegerValue())
            .collect();
        if shape.len() != strides.len() {
            return Err(Error::Invalid(format!(
                "Invalid tensor stride rank: shape={shape:?}, strides={strides:?}"
            )));
        }
        let maximum = shape
            .iter()
            .zip(&strides)
            .try_fold(0usize, |sum, (&dimension, &stride)| {
                sum.checked_add((dimension - 1).checked_mul(stride)?)
            });
        let minimum_bytes = maximum
            .and_then(|n| n.checked_add(1))
            .and_then(|n| n.checked_mul(std::mem::size_of::<T>()))
            .ok_or_else(|| {
                Error::Invalid(format!(
                    "Tensor backing size overflow: shape={shape:?}, strides={strides:?}"
                ))
            })?;
        Ok(Self {
            shape,
            strides,
            count,
            minimum_bytes,
        })
    }

    fn check_bytes(&self, size: isize) -> Result<()> {
        if size < 0 || self.minimum_bytes > size as usize {
            return Err(Error::Invalid(format!(
                "Invalid Core ML backing: shape={:?}, strides={:?}, bytes={size}, required={}",
                self.shape, self.strides, self.minimum_bytes
            )));
        }
        Ok(())
    }

    fn check_array<T: Scalar>(&self, array: &MLMultiArray) -> Result<()> {
        // SAFETY: This reads metadata only; no raw pointer escapes.
        let shape = unsafe { array.shape() };
        let strides = unsafe { array.strides() };
        if unsafe { array.dataType() } != T::TYPE
            || shape.len() != self.shape.len()
            || strides.len() != self.strides.len()
            || shape
                .iter()
                .zip(&self.shape)
                .any(|(a, &b)| a.unsignedIntegerValue() != b)
            || strides
                .iter()
                .zip(&self.strides)
                .any(|(a, &b)| a.unsignedIntegerValue() != b)
        {
            return Err(Error::Invalid(format!(
                "Tensor contract changed: expected shape={:?}, strides={:?}, dtype={:?}; actual shape={:?}, dtype={:?}",
                self.shape,
                self.strides,
                T::TYPE,
                array_shape(array),
                unsafe { array.dataType() }
            )));
        }
        Ok(())
    }
}

/// Owns one input or output backing across prediction steps.
pub(super) struct TensorBuffer<T: Scalar> {
    array: Retained<MLMultiArray>,
    layout: TensorLayout,
    scalar: std::marker::PhantomData<T>,
}

impl<T: Scalar> TensorBuffer<T> {
    pub(super) fn from_slice(shape: &[usize], data: &[T]) -> Result<Self> {
        if shape_product(shape)? != data.len() {
            return Err(Error::Invalid(format!(
                "Tensor size mismatch: shape={shape:?}, elements={}",
                data.len()
            )));
        }
        let dimensions = NSArray::from_retained_slice(
            &shape
                .iter()
                .map(|&d| NSNumber::new_usize(d))
                .collect::<Vec<_>>(),
        );
        // SAFETY: Dimensions are positive and T declares the exact Core ML dtype.
        let array = unsafe {
            MLMultiArray::initWithShape_dataType_error(MLMultiArray::alloc(), &dimensions, T::TYPE)
        }
        .map_err(|e| core_error("allocate multi-array", e))?;
        let layout = TensorLayout::new::<T>(&array)?;
        let mut buffer = Self {
            array,
            layout,
            scalar: std::marker::PhantomData,
        };
        buffer.replace(data.iter().copied())?;
        Ok(buffer)
    }

    pub(super) fn zero(shape: &[usize]) -> Result<Self> {
        Self::from_slice(shape, &vec![T::default(); shape_product(shape)?])
    }

    pub(super) fn array(&self) -> &MLMultiArray {
        &self.array
    }

    pub(super) fn into_array(self) -> Retained<MLMultiArray> {
        self.array
    }

    pub(super) fn replace<I: ExactSizeIterator<Item = T>>(&mut self, values: I) -> Result<()> {
        if values.len() != self.layout.count {
            return Err(Error::Invalid(format!(
                "Reusable tensor length mismatch: shape={:?}, expected={}, actual={}",
                self.layout.shape,
                self.layout.count,
                values.len()
            )));
        }
        let values = RefCell::new(Some(values));
        let status = RefCell::new(None);
        let layout = &self.layout;
        let handler = StackBlock::new(
            |pointer: NonNull<c_void>, size: isize, strides: NonNull<NSArray<NSNumber>>| {
                let result = (|| {
                    layout.check_bytes(size)?;
                    // SAFETY: The callback supplies metadata valid during this callback.
                    let actual = unsafe { strides.as_ref() };
                    if actual.len() != layout.strides.len()
                        || actual
                            .iter()
                            .zip(&layout.strides)
                            .any(|(a, &b)| a.unsignedIntegerValue() != b)
                    {
                        return Err(Error::Invalid(
                            "Reusable tensor writable strides changed".into(),
                        ));
                    }
                    let values = values.borrow_mut().take().ok_or_else(|| {
                        Error::Invalid("Core ML invoked tensor writer more than once".into())
                    })?;
                    for (i, value) in values.enumerate() {
                        // SAFETY: Layout and byte bounds were validated; the dtype is T.
                        unsafe {
                            pointer
                                .cast::<T>()
                                .as_ptr()
                                .add(flat_offset(i, &layout.shape, &layout.strides))
                                .write(value)
                        };
                    }
                    Ok(())
                })();
                *status.borrow_mut() = Some(result);
            },
        );
        // SAFETY: This stack block is synchronous and its raw pointer never escapes.
        unsafe { self.array.getMutableBytesWithHandler(&handler) };
        status
            .into_inner()
            .ok_or_else(|| Error::Invalid("Core ML did not provide mutable tensor bytes".into()))?
    }
}

/// Refreshes a host snapshot without allocating a new Vec for every prediction.
pub(super) struct TensorReader<T: Scalar> {
    layout: TensorLayout,
    values: Vec<T>,
}

impl<T: Scalar> TensorReader<T> {
    pub(super) fn new(array: &MLMultiArray) -> Result<Self> {
        let layout = TensorLayout::new::<T>(array)?;
        let values = vec![T::default(); layout.count];
        Ok(Self { layout, values })
    }

    pub(super) fn refresh(&mut self, array: &MLMultiArray) -> Result<()> {
        self.layout.check_array::<T>(array)?;
        let layout = &self.layout;
        let values = RefCell::new(&mut self.values);
        let status = RefCell::new(None);
        let handler = StackBlock::new(|pointer: NonNull<c_void>, size: isize| {
            let result = (|| {
                layout.check_bytes(size)?;
                let mut values = values.borrow_mut();
                for (i, value) in values.iter_mut().enumerate() {
                    // SAFETY: T, strided offsets and backing bounds were checked.
                    *value = unsafe {
                        pointer
                            .cast::<T>()
                            .as_ptr()
                            .add(flat_offset(i, &layout.shape, &layout.strides))
                            .read()
                    };
                }
                Ok(())
            })();
            *status.borrow_mut() = Some(result);
        });
        // SAFETY: The read callback is synchronous and does not retain the pointer.
        unsafe { array.getBytesWithHandler(&handler) };
        status
            .into_inner()
            .ok_or_else(|| Error::Invalid("Core ML did not provide tensor bytes".into()))?
    }

    pub(super) fn values(&self) -> &[T] {
        &self.values
    }
}

pub(super) fn multi_array<T: Scalar>(
    shape: &[usize],
    data: &[T],
) -> Result<Retained<MLMultiArray>> {
    Ok(TensorBuffer::from_slice(shape, data)?.into_array())
}

pub(super) fn read_array<T: Scalar>(array: &MLMultiArray) -> Result<Vec<T>> {
    let mut reader = TensorReader::new(array)?;
    reader.refresh(array)?;
    Ok(reader.values)
}

pub(super) fn read_floating_tensor(array: &MLMultiArray) -> Result<Vec<f32>> {
    // SAFETY: Inspecting the declared dtype does not access raw storage.
    match unsafe { array.dataType() } {
        MLMultiArrayDataType::Float32 => read_array::<f32>(array),
        MLMultiArrayDataType::Float16 => Ok(read_array::<half::f16>(array)?
            .into_iter()
            .map(half::f16::to_f32)
            .collect()),
        dtype => Err(Error::Invalid(format!(
            "Unsupported Core ML floating tensor dtype: {dtype:?}"
        ))),
    }
}

#[cfg(test)]
mod reuse_tests {
    use super::*;

    #[test]
    fn reusable_input_keeps_its_array_and_replaces_values() {
        autoreleasepool(|_| {
            let mut buffer = TensorBuffer::<f32>::zero(&[1, 3]).unwrap();
            let original = buffer.array() as *const MLMultiArray;
            buffer.replace([1.0, 2.0, 3.0].into_iter()).unwrap();
            assert_eq!(read_array::<f32>(buffer.array()).unwrap(), [1.0, 2.0, 3.0]);
            buffer.replace([4.0, 5.0, 6.0].into_iter()).unwrap();
            assert_eq!(buffer.array() as *const MLMultiArray, original);
            assert_eq!(read_array::<f32>(buffer.array()).unwrap(), [4.0, 5.0, 6.0]);
            assert!(buffer.replace([7.0, 8.0].into_iter()).is_err());
            assert_eq!(read_array::<f32>(buffer.array()).unwrap(), [4.0, 5.0, 6.0]);
        });
    }

    #[test]
    fn reusable_reader_refreshes_without_replacing_its_storage() {
        autoreleasepool(|_| {
            let mut buffer = TensorBuffer::<f32>::zero(&[1, 3]).unwrap();
            let mut reader = TensorReader::<f32>::new(buffer.array()).unwrap();
            let original = reader.values().as_ptr();
            buffer.replace([1.0, 2.0, 3.0].into_iter()).unwrap();
            reader.refresh(buffer.array()).unwrap();
            assert_eq!(reader.values(), [1.0, 2.0, 3.0]);
            buffer.replace([4.0, 5.0, 6.0].into_iter()).unwrap();
            reader.refresh(buffer.array()).unwrap();
            assert_eq!(reader.values().as_ptr(), original);
            assert_eq!(reader.values(), [4.0, 5.0, 6.0]);
            let wrong = multi_array(&[3, 1], &[1.0f32, 2.0, 3.0]).unwrap();
            assert!(reader.refresh(&wrong).is_err());
            assert_eq!(reader.values(), [4.0, 5.0, 6.0]);
        });
    }
}