tflite_c/

tensor.rs

//! Borrowed views into a [`crate::Interpreter`]'s input and output tensors.

use std::ffi::CStr;
use std::marker::PhantomData;
use std::os::raw::c_int;
use std::sync::Arc;

use crate::error::{Error, Result};
use crate::ffi::{TfLiteQuantizationParams, TfLiteTensor, TfLiteType};
use crate::library::TfLiteLibrary;

/// Read-only borrow of a tensor owned by a [`crate::Interpreter`].
///
/// `Tensor<'a>` cannot outlive the interpreter it came from — the borrow
/// checker enforces this by attaching a phantom lifetime to the underlying
/// pointer.
pub struct Tensor<'a> {
    raw: *const TfLiteTensor,
    lib: Arc<TfLiteLibrary>,
    _marker: PhantomData<&'a ()>,
}

/// Read-write borrow of a tensor owned by a [`crate::Interpreter`].
///
/// Like [`Tensor`] but additionally exposes [`TensorMut::data_mut`] for
/// filling input buffers in place.
pub struct TensorMut<'a> {
    raw: *mut TfLiteTensor,
    lib: Arc<TfLiteLibrary>,
    _marker: PhantomData<&'a mut ()>,
}

impl<'a> Tensor<'a> {
    /// Construct from a raw pointer; intended for use by the interpreter only.
    ///
    /// # Safety
    /// The caller must guarantee that `raw` is a valid `TfLiteTensor*` that
    /// outlives `'a` and is not being mutated concurrently.
    pub(crate) unsafe fn from_raw(raw: *const TfLiteTensor, lib: Arc<TfLiteLibrary>) -> Self {
        Self { raw, lib, _marker: PhantomData }
    }

    /// Tensor shape as a list of per-axis dimensions.
    pub fn dims(&self) -> Vec<i32> {
        common_dims(self.raw, &self.lib)
    }

    /// Element type of the tensor.
    pub fn dtype(&self) -> TfLiteType {
        common_dtype(self.raw, &self.lib)
    }

    /// Tensor name as reported by the model, if any.
    pub fn name(&self) -> Option<&str> {
        common_name(self.raw, &self.lib)
    }

    /// Total size of the tensor's data buffer in bytes.
    pub fn byte_size(&self) -> usize {
        common_byte_size(self.raw, &self.lib)
    }

    /// Affine quantization parameters (`scale`, `zero_point`).
    pub fn quantization(&self) -> TfLiteQuantizationParams {
        common_quantization(self.raw, &self.lib)
    }

    /// Raw byte view of the tensor's data buffer.
    ///
    /// Returns an empty slice if the C API reports a NULL data pointer.
    pub fn data(&self) -> &[u8] {
        common_data(self.raw, &self.lib)
    }

    /// Typed `&[f32]` view; errors if the tensor is not `Float32`.
    pub fn as_slice_f32(&self) -> Result<&[f32]> {
        typed_slice::<f32>(self.raw, &self.lib, TfLiteType::Float32)
    }

    /// Typed `&[i8]` view; errors if the tensor is not `Int8`.
    pub fn as_slice_i8(&self) -> Result<&[i8]> {
        typed_slice::<i8>(self.raw, &self.lib, TfLiteType::Int8)
    }

    /// Typed `&[u8]` view; errors if the tensor is not `UInt8`.
    pub fn as_slice_u8(&self) -> Result<&[u8]> {
        typed_slice::<u8>(self.raw, &self.lib, TfLiteType::UInt8)
    }

    /// Typed `&[i32]` view; errors if the tensor is not `Int32`.
    pub fn as_slice_i32(&self) -> Result<&[i32]> {
        typed_slice::<i32>(self.raw, &self.lib, TfLiteType::Int32)
    }

    /// Read the tensor as a fresh `Vec<f32>`, dequantizing `Int8`/`UInt8`
    /// using the tensor's quantization parameters.
    ///
    /// The dequantization formula is `real = (q - zero_point) * scale`.
    /// `Float32` tensors are copied straight through.
    pub fn to_vec_f32(&self) -> Result<Vec<f32>> {
        match self.dtype() {
            TfLiteType::Float32 => self.as_slice_f32().map(<[f32]>::to_vec),
            TfLiteType::Int8 => {
                let raw = self.as_slice_i8()?;
                Ok(dequantize_i8(raw, self.quantization()))
            }
            TfLiteType::UInt8 => {
                let raw = self.as_slice_u8()?;
                Ok(dequantize_u8(raw, self.quantization()))
            }
            actual => Err(Error::DtypeMismatch { expected: TfLiteType::Float32, actual }),
        }
    }
}

impl<'a> TensorMut<'a> {
    /// Construct from a raw pointer; intended for use by the interpreter only.
    ///
    /// # Safety
    /// The caller must guarantee that `raw` is a valid `TfLiteTensor*` that
    /// outlives `'a` and that the unique-borrow invariant is upheld.
    pub(crate) unsafe fn from_raw(raw: *mut TfLiteTensor, lib: Arc<TfLiteLibrary>) -> Self {
        Self { raw, lib, _marker: PhantomData }
    }

    /// Tensor shape as a list of per-axis dimensions.
    pub fn dims(&self) -> Vec<i32> {
        common_dims(self.raw, &self.lib)
    }

    /// Element type of the tensor.
    pub fn dtype(&self) -> TfLiteType {
        common_dtype(self.raw, &self.lib)
    }

    /// Tensor name as reported by the model, if any.
    pub fn name(&self) -> Option<&str> {
        common_name(self.raw, &self.lib)
    }

    /// Total size of the tensor's data buffer in bytes.
    pub fn byte_size(&self) -> usize {
        common_byte_size(self.raw, &self.lib)
    }

    /// Affine quantization parameters (`scale`, `zero_point`).
    pub fn quantization(&self) -> TfLiteQuantizationParams {
        common_quantization(self.raw, &self.lib)
    }

    /// Read-only byte view of the tensor's data buffer.
    pub fn data(&self) -> &[u8] {
        common_data(self.raw, &self.lib)
    }

    /// Mutable byte view of the tensor's data buffer.
    ///
    /// Errors with [`Error::NullData`] if the C API reports a NULL data
    /// pointer (which can happen if the tensor has not been allocated yet).
    pub fn data_mut(&mut self) -> Result<&mut [u8]> {
        // SAFETY: `raw` is non-null per the interpreter's bounds check, and
        // `TensorMut` requires exclusive access to the tensor.
        unsafe {
            let ptr = (self.lib.tensor_data)(self.raw) as *mut u8;
            if ptr.is_null() {
                return Err(Error::NullData);
            }
            let len = (self.lib.tensor_byte_size)(self.raw);
            Ok(std::slice::from_raw_parts_mut(ptr, len))
        }
    }
}

// --- Shared helpers (deliberately free functions to avoid trait dispatch) -------

fn common_dims(raw: *const TfLiteTensor, lib: &TfLiteLibrary) -> Vec<i32> {
    // SAFETY: `raw` is non-null per the interpreter's bounds check.
    unsafe {
        let n = (lib.tensor_num_dims)(raw);
        if n <= 0 {
            return Vec::new();
        }
        (0..n).map(|i| (lib.tensor_dim)(raw, i as c_int)).collect()
    }
}

fn common_dtype(raw: *const TfLiteTensor, lib: &TfLiteLibrary) -> TfLiteType {
    // SAFETY: `raw` is non-null per the interpreter's bounds check.
    unsafe { (lib.tensor_type)(raw) }
}

fn common_name<'a>(raw: *const TfLiteTensor, lib: &TfLiteLibrary) -> Option<&'a str> {
    // SAFETY: `raw` is non-null per the interpreter's bounds check; the
    // returned C string is owned by TFLite and lives at least as long as
    // the tensor handle, which is upheld by the borrow checker via the
    // caller's `'a` lifetime.
    unsafe {
        let ptr = (lib.tensor_name)(raw);
        if ptr.is_null() {
            return None;
        }
        CStr::from_ptr(ptr).to_str().ok()
    }
}

fn common_byte_size(raw: *const TfLiteTensor, lib: &TfLiteLibrary) -> usize {
    // SAFETY: `raw` is non-null per the interpreter's bounds check.
    unsafe { (lib.tensor_byte_size)(raw) }
}

fn common_quantization(
    raw: *const TfLiteTensor,
    lib: &TfLiteLibrary,
) -> TfLiteQuantizationParams {
    // SAFETY: `raw` is non-null per the interpreter's bounds check.
    unsafe { (lib.tensor_quantization_params)(raw) }
}

fn common_data<'a>(raw: *const TfLiteTensor, lib: &TfLiteLibrary) -> &'a [u8] {
    // SAFETY: `raw` is non-null per the interpreter's bounds check.
    unsafe {
        let ptr = (lib.tensor_data)(raw) as *const u8;
        if ptr.is_null() {
            return &[];
        }
        let len = (lib.tensor_byte_size)(raw);
        std::slice::from_raw_parts(ptr, len)
    }
}

fn typed_slice<'a, T: Copy>(
    raw: *const TfLiteTensor,
    lib: &TfLiteLibrary,
    expected: TfLiteType,
) -> Result<&'a [T]> {
    let actual = common_dtype(raw, lib);
    if actual != expected {
        return Err(Error::DtypeMismatch { expected, actual });
    }
    // SAFETY: We have just confirmed the tensor's element type matches `T`'s
    // representation. `byte_size` reports the total data length in bytes.
    unsafe {
        let ptr = (lib.tensor_data)(raw) as *const T;
        if ptr.is_null() {
            return Err(Error::NullData);
        }
        let bytes = (lib.tensor_byte_size)(raw);
        let len = bytes / std::mem::size_of::<T>();
        Ok(std::slice::from_raw_parts(ptr, len))
    }
}

/// Dequantize an int8 buffer into f32 using affine quantization.
///
/// `real = (q - zero_point) * scale`. Exposed at crate-private scope so it
/// can be unit-tested without a live TFLite library.
pub(crate) fn dequantize_i8(raw: &[i8], params: TfLiteQuantizationParams) -> Vec<f32> {
    let zp = params.zero_point as f32;
    let scale = params.scale;
    raw.iter().map(|&q| (q as f32 - zp) * scale).collect()
}

/// Dequantize a uint8 buffer into f32 using affine quantization.
pub(crate) fn dequantize_u8(raw: &[u8], params: TfLiteQuantizationParams) -> Vec<f32> {
    let zp = params.zero_point as f32;
    let scale = params.scale;
    raw.iter().map(|&q| (q as f32 - zp) * scale).collect()
}

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

    #[test]
    fn dequantize_i8_matches_formula() {
        // scale=0.5, zero_point=-10 → real = (q + 10) * 0.5
        let params = TfLiteQuantizationParams { scale: 0.5, zero_point: -10 };
        let raw: Vec<i8> = vec![-10, 0, 10, 20];
        let got = dequantize_i8(&raw, params);
        assert_eq!(got, vec![0.0, 5.0, 10.0, 15.0]);
    }

    #[test]
    fn dequantize_i8_with_zero_scale_is_all_zero() {
        let params = TfLiteQuantizationParams { scale: 0.0, zero_point: 42 };
        let got = dequantize_i8(&[1, 2, 3, 4], params);
        assert_eq!(got, vec![0.0, 0.0, 0.0, 0.0]);
    }

    #[test]
    fn dequantize_u8_matches_formula() {
        // scale=0.25, zero_point=128 → real = (q - 128) * 0.25
        let params = TfLiteQuantizationParams { scale: 0.25, zero_point: 128 };
        let raw: Vec<u8> = vec![128, 132, 124, 0];
        let got = dequantize_u8(&raw, params);
        assert_eq!(got, vec![0.0, 1.0, -1.0, -32.0]);
    }

    #[test]
    fn dequantize_handles_empty_input() {
        let params = TfLiteQuantizationParams { scale: 1.0, zero_point: 0 };
        assert!(dequantize_i8(&[], params).is_empty());
        assert!(dequantize_u8(&[], params).is_empty());
    }
}