tflite-c-rs 0.0.1

//! Safe wrappers around `TfLiteInterpreterOptions` and `TfLiteInterpreter`.

use std::os::raw::c_int;
use std::sync::Arc;

use crate::error::{Error, Result};
use crate::ffi::{TfLiteInterpreter, TfLiteInterpreterOptions, TfLiteStatus};
use crate::library::{ExternalDelegate, TfLiteLibrary};
use crate::model::Model;
use crate::tensor::{Tensor, TensorMut};

/// Builder for an interpreter's configuration (thread count, delegates, …).
///
/// Each setter returns `&mut Self` so the builder reads naturally inline:
/// `InterpreterOptions::new(lib).num_threads(2).add_delegate(d)`.
pub struct InterpreterOptions {
    raw: *mut TfLiteInterpreterOptions,
    lib: Arc<TfLiteLibrary>,
}

// SAFETY: `InterpreterOptions` owns a raw C handle; access is exclusive
// (it is consumed by `Interpreter::new`), so cross-thread aliasing cannot
// happen in safe code.
unsafe impl Send for InterpreterOptions {}
unsafe impl Sync for InterpreterOptions {}

impl InterpreterOptions {
    /// Allocate a fresh `TfLiteInterpreterOptions` and return a builder.
    pub fn new(lib: Arc<TfLiteLibrary>) -> Self {
        // SAFETY: `options_create` takes no arguments and is safe to call.
        let raw = unsafe { (lib.options_create)() };
        Self { raw, lib }
    }

    /// Number of CPU threads the interpreter may use during invoke.
    pub fn num_threads(&mut self, n: i32) -> &mut Self {
        // SAFETY: `raw` is non-null (created by `options_create`).
        unsafe {
            (self.lib.options_set_num_threads)(self.raw, n as c_int);
        }
        self
    }

    /// Attach an [`ExternalDelegate`] (e.g. NNAPI, GPU, XNNPACK, VX).
    ///
    /// TensorFlow Lite takes ownership of the delegate handle from this
    /// point forward: it will release the handle when the interpreter is
    /// destroyed.
    pub fn add_delegate(&mut self, delegate: ExternalDelegate) -> &mut Self {
        // SAFETY: `raw` is non-null (created by `options_create`), and the
        // delegate handle was just produced by `TfLiteExternalDelegateCreate`.
        unsafe {
            (self.lib.options_add_delegate)(self.raw, delegate.as_ptr());
        }

        // SAFETY: Once `TfLiteInterpreterOptionsAddDelegate` has been
        // called, TFLite owns the delegate and will free it when the
        // interpreter is destroyed. Running `ExternalDelegate::drop` here
        // would call `TfLiteExternalDelegateDelete` on the same pointer,
        // producing a double free. `mem::forget` skips the destructor
        // while still releasing the `Arc<TfLiteLibrary>` that the delegate
        // was holding (since the `Arc` is moved out of `delegate` by
        // value — but `mem::forget` leaks it instead of decrementing).
        //
        // The leaked `Arc` count is bounded by the number of delegates a
        // user attaches in the lifetime of the process; for the typical
        // "one interpreter per process" pattern this is at most a handful.
        std::mem::forget(delegate);
        self
    }

    /// Returns the raw `TfLiteInterpreterOptions*`. Advanced use only.
    pub fn as_ptr(&self) -> *const TfLiteInterpreterOptions {
        self.raw
    }
}

impl Drop for InterpreterOptions {
    fn drop(&mut self) {
        if !self.raw.is_null() {
            // SAFETY: `raw` was allocated by `options_create` and is not
            // aliased — `Interpreter::new` does not call `options_delete`
            // itself; TFLite's contract is that the caller deletes options
            // after the interpreter is built.
            unsafe {
                (self.lib.options_delete)(self.raw);
            }
        }
    }
}

/// A constructed TFLite interpreter, ready to invoke.
///
/// Construction calls `TfLiteInterpreterCreate` and then
/// `TfLiteInterpreterAllocateTensors`, so tensor buffers are valid as soon
/// as `Interpreter::new` returns `Ok(...)`.
///
/// `Interpreter` owns its [`Model`] and [`InterpreterOptions`]; both are
/// freed (in the correct order, after the interpreter itself) when the
/// interpreter is dropped.
pub struct Interpreter {
    raw: *mut TfLiteInterpreter,
    lib: Arc<TfLiteLibrary>,
    _model: Model,
    _options: InterpreterOptions,
}

// SAFETY: The TFLite C interpreter is not internally synchronized — calls
// into a single interpreter from multiple threads simultaneously are
// undefined behaviour. We mark `Send`/`Sync` so the type can live inside
// the canonical `Arc<Mutex<Interpreter>>` pattern; downstream code must
// serialize calls itself.
unsafe impl Send for Interpreter {}
unsafe impl Sync for Interpreter {}

impl Interpreter {
    /// Build a new interpreter and allocate its tensors.
    pub fn new(model: Model, options: InterpreterOptions) -> Result<Self> {
        let lib = model.lib().clone();

        // SAFETY: Both handles were created by the matching TFLite C calls
        // and are non-null. `interpreter_create` borrows model and options
        // for its own bookkeeping; the wrapper keeps both alive in the
        // returned struct.
        let raw = unsafe { (lib.interpreter_create)(model.as_ptr(), options.as_ptr()) };
        if raw.is_null() {
            return Err(Error::InterpreterCreateFailed);
        }

        // SAFETY: `raw` is freshly returned and non-null.
        let status = unsafe { (lib.interpreter_allocate_tensors)(raw) };
        if status != TfLiteStatus::Ok {
            // SAFETY: We own `raw` exclusively; clean it up before erroring.
            unsafe {
                (lib.interpreter_delete)(raw);
            }
            return Err(Error::AllocateTensorsFailed);
        }

        Ok(Self { raw, lib, _model: model, _options: options })
    }

    /// Number of input tensors exposed by the model.
    pub fn input_count(&self) -> usize {
        // SAFETY: `raw` is non-null for the lifetime of `self`.
        unsafe { (self.lib.interpreter_get_input_tensor_count)(self.raw) as usize }
    }

    /// Number of output tensors exposed by the model.
    pub fn output_count(&self) -> usize {
        // SAFETY: `raw` is non-null for the lifetime of `self`.
        unsafe { (self.lib.interpreter_get_output_tensor_count)(self.raw) as usize }
    }

    /// Borrow input tensor `index` immutably.
    pub fn input(&self, index: usize) -> Result<Tensor<'_>> {
        let count = self.input_count();
        if index >= count {
            return Err(Error::TensorIndexOutOfBounds { index, count });
        }
        // SAFETY: `index` is in bounds; the returned pointer is borrowed
        // from `self` and lives at least as long as `'_`.
        unsafe {
            let raw = (self.lib.interpreter_get_input_tensor)(self.raw, index as c_int);
            if raw.is_null() {
                return Err(Error::NullTensor);
            }
            Ok(Tensor::from_raw(raw as *const _, self.lib.clone()))
        }
    }

    /// Borrow input tensor `index` mutably (for filling input buffers).
    pub fn input_mut(&mut self, index: usize) -> Result<TensorMut<'_>> {
        let count = self.input_count();
        if index >= count {
            return Err(Error::TensorIndexOutOfBounds { index, count });
        }
        // SAFETY: `&mut self` guarantees exclusive access; the returned
        // pointer lives at least as long as `'_`.
        unsafe {
            let raw = (self.lib.interpreter_get_input_tensor)(self.raw, index as c_int);
            if raw.is_null() {
                return Err(Error::NullTensor);
            }
            Ok(TensorMut::from_raw(raw, self.lib.clone()))
        }
    }

    /// Borrow output tensor `index` immutably.
    pub fn output(&self, index: usize) -> Result<Tensor<'_>> {
        let count = self.output_count();
        if index >= count {
            return Err(Error::TensorIndexOutOfBounds { index, count });
        }
        // SAFETY: `index` is in bounds; the returned pointer is borrowed
        // from `self` and lives at least as long as `'_`.
        unsafe {
            let raw = (self.lib.interpreter_get_output_tensor)(self.raw, index as c_int);
            if raw.is_null() {
                return Err(Error::NullTensor);
            }
            Ok(Tensor::from_raw(raw, self.lib.clone()))
        }
    }

    /// Run inference. Maps non-`Ok` statuses to [`Error::InvokeFailed`].
    pub fn invoke(&mut self) -> Result<()> {
        // SAFETY: `&mut self` guarantees exclusive access; `raw` is non-null.
        let status = unsafe { (self.lib.interpreter_invoke)(self.raw) };
        if status == TfLiteStatus::Ok {
            Ok(())
        } else {
            Err(Error::InvokeFailed { status })
        }
    }
}

impl Drop for Interpreter {
    fn drop(&mut self) {
        // SAFETY: `raw` was returned by `interpreter_create`. We own it
        // exclusively, and the model and options live in `_model` /
        // `_options` so they are dropped *after* the interpreter — which
        // matches TFLite's required destruction order.
        unsafe {
            (self.lib.interpreter_delete)(self.raw);
        }
    }
}