tflite-c-rs 0.0.1

//! Runtime loader for the TensorFlow Lite C shared library.
//!
//! This module is the only place in the crate that touches `libloading`.
//! Once a [`TfLiteLibrary`] is constructed, every other type just calls into
//! the cached function pointers it owns.

use std::ffi::CString;
use std::path::Path;
use std::sync::Arc;

use libloading::{Library, Symbol};

use crate::error::{Error, Result};
use crate::ffi::*;

/// Default candidate paths probed by [`TfLiteLibrary::load_default`].
///
/// The list covers the common shared-library names used by upstream and
/// distro packages on Linux, macOS, and Windows. Names without a directory
/// component fall back to the platform loader's standard search path
/// (`LD_LIBRARY_PATH` on Linux, etc.).
pub const DEFAULT_LIBRARY_PATHS: &[&str] = &[
    // Bare names — picked up via the loader's search path.
    "libtensorflowlite_c.so",
    "libtensorflow-lite.so",
    "libtensorflowlite_c.dylib",
    "libtensorflow-lite.dylib",
    "tensorflowlite_c.dll",
    // Common Linux install prefixes.
    "/usr/lib/libtensorflowlite_c.so",
    "/usr/lib/libtensorflow-lite.so",
    "/usr/local/lib/libtensorflowlite_c.so",
    "/usr/local/lib/libtensorflow-lite.so",
    // Common macOS install prefixes (system + Homebrew, Intel and Apple Silicon).
    "/usr/local/lib/libtensorflowlite_c.dylib",
    "/opt/homebrew/lib/libtensorflowlite_c.dylib",
];

/// Resolved handle to a dynamically loaded TensorFlow Lite C library.
///
/// Construct via [`TfLiteLibrary::load_default`] (probes a set of common
/// install paths) or [`TfLiteLibrary::load_from_path`] (explicit path). The
/// returned [`Arc`] is cloned into every wrapper type so that the underlying
/// shared object remains mapped for as long as anything still references it.
pub struct TfLiteLibrary {
    // The underlying loader handle. Kept alive for the process; cached
    // function pointers below borrow from a leaked clone of this same
    // `Arc<Library>`, so dropping `TfLiteLibrary` does not invalidate them.
    _lib: Arc<Library>,

    // --- Core model / interpreter symbols ------------------------------------
    pub(crate) model_create_from_file: Symbol<'static, ModelCreateFromFileFn>,
    pub(crate) model_delete: Symbol<'static, ModelDeleteFn>,

    pub(crate) options_create: Symbol<'static, OptionsCreateFn>,
    pub(crate) options_delete: Symbol<'static, OptionsDeleteFn>,
    pub(crate) options_set_num_threads: Symbol<'static, OptionsSetNumThreadsFn>,
    pub(crate) options_add_delegate: Symbol<'static, OptionsAddDelegateFn>,

    pub(crate) interpreter_create: Symbol<'static, InterpreterCreateFn>,
    pub(crate) interpreter_delete: Symbol<'static, InterpreterDeleteFn>,
    pub(crate) interpreter_allocate_tensors: Symbol<'static, InterpreterAllocateTensorsFn>,
    pub(crate) interpreter_invoke: Symbol<'static, InterpreterInvokeFn>,
    pub(crate) interpreter_get_input_tensor_count: Symbol<'static, InterpreterGetTensorCountFn>,
    pub(crate) interpreter_get_output_tensor_count: Symbol<'static, InterpreterGetTensorCountFn>,
    pub(crate) interpreter_get_input_tensor: Symbol<'static, InterpreterGetInputTensorFn>,
    pub(crate) interpreter_get_output_tensor: Symbol<'static, InterpreterGetOutputTensorFn>,

    // --- Tensor introspection ----------------------------------------------
    pub(crate) tensor_type: Symbol<'static, TensorTypeFn>,
    pub(crate) tensor_num_dims: Symbol<'static, TensorNumDimsFn>,
    pub(crate) tensor_dim: Symbol<'static, TensorDimFn>,
    pub(crate) tensor_byte_size: Symbol<'static, TensorByteSizeFn>,
    pub(crate) tensor_data: Symbol<'static, TensorDataFn>,
    pub(crate) tensor_name: Symbol<'static, TensorNameFn>,
    pub(crate) tensor_quantization_params: Symbol<'static, TensorQuantizationParamsFn>,

    // --- External delegate API (optional — TFLite's c_api_experimental.h) ---
    //
    // These symbols come from the *experimental* header and are routinely
    // omitted from prebuilt TFLite C shared objects (e.g. community builds,
    // some vendor SDKs). Resolved lazily as a group: present if and only if
    // every one of the six symbols loaded. `ExternalDelegate::builder`
    // returns [`Error::ExternalDelegateApiUnavailable`] when the group is
    // missing rather than failing the entire library load.
    pub(crate) external_delegate: Option<ExternalDelegateSymbols>,
}

pub(crate) struct ExternalDelegateSymbols {
    pub options_create: Symbol<'static, ExternalDelegateOptionsCreateFn>,
    pub options_delete: Symbol<'static, ExternalDelegateOptionsDeleteFn>,
    pub options_set_library_path: Symbol<'static, ExternalDelegateOptionsSetLibraryPathFn>,
    pub options_insert: Symbol<'static, ExternalDelegateOptionsInsertFn>,
    pub create: Symbol<'static, ExternalDelegateCreateFn>,
    pub delete: Symbol<'static, ExternalDelegateDeleteFn>,
}

// SAFETY: `TfLiteLibrary` owns an `Arc<Library>` and a fixed set of cached
// function pointers. The function pointers are immutable after `load_*`
// returns, and `libloading::Library` is itself `Send + Sync`. Sharing a
// `&TfLiteLibrary` across threads therefore just lets each thread call
// `extern "C"` functions whose thread-safety is governed by TFLite's own
// contract (which the wrapper types document separately).
unsafe impl Send for TfLiteLibrary {}
unsafe impl Sync for TfLiteLibrary {}

impl TfLiteLibrary {
    /// Probe a default list of well-known paths and return the first library
    /// that loads successfully.
    ///
    /// See [`DEFAULT_LIBRARY_PATHS`] for the exact set of candidates. Returns
    /// [`Error::LoadFailed`] (with the full probe list) if none open.
    pub fn load_default() -> Result<Arc<Self>> {
        let mut last_err: Option<libloading::Error> = None;
        for candidate in DEFAULT_LIBRARY_PATHS {
            match Self::try_load(candidate) {
                Ok(lib) => return Ok(lib),
                Err(Error::LoadFailed { source, .. }) => last_err = Some(source),
                Err(other) => return Err(other),
            }
        }
        Err(Error::LoadFailed {
            source: last_err.unwrap_or_else(|| {
                // Synthesize an error if the candidate list was empty —
                // should be impossible but keeps the API total.
                libloading::Error::DlOpenUnknown
            }),
            paths: DEFAULT_LIBRARY_PATHS.iter().map(|s| s.to_string()).collect(),
        })
    }

    /// Load the TensorFlow Lite shared library from an explicit filesystem path.
    pub fn load_from_path<P: AsRef<Path>>(path: P) -> Result<Arc<Self>> {
        let path_ref = path.as_ref();
        Self::try_load(path_ref).map_err(|err| match err {
            Error::LoadFailed { source, .. } => Error::LoadFailed {
                source,
                paths: vec![path_ref.display().to_string()],
            },
            other => other,
        })
    }

    fn try_load<P: AsRef<Path>>(path: P) -> Result<Arc<Self>> {
        let path_ref = path.as_ref();

        // SAFETY: `Library::new` runs `dlopen`/`LoadLibrary`, which executes
        // the target shared object's initializers. We trust the caller (or
        // the default probe list) to point at a real TFLite build.
        let lib = unsafe { Library::new(path_ref) }.map_err(|source| Error::LoadFailed {
            source,
            paths: vec![path_ref.display().to_string()],
        })?;
        let lib = Arc::new(lib);

        // SAFETY: `libloading::Symbol` borrows from a `Library`, and the
        // `Symbol`s stored in this struct need a `'static` lifetime so the
        // struct itself can be `'static`-friendly (e.g. wrapped in `Arc`
        // and passed across threads).
        //
        // We obtain `'static` by `Box::leak`ing a clone of the `Arc<Library>`.
        // The leak is bounded: at most one allocation per successful
        // `load_default` / `load_from_path` call, and applications typically
        // load TFLite once for the lifetime of the process. The original
        // `Arc<Library>` is also kept inside the returned struct (`_lib`),
        // so the underlying handle is reference-counted and not freed early.
        let leaked: &'static Arc<Library> = Box::leak(Box::new(Arc::clone(&lib)));

        macro_rules! sym {
            ($name:literal) => {{
                // SAFETY: `get` looks up a C symbol by name. We use the
                // function-pointer types declared in `ffi.rs`, which match
                // the upstream signatures byte-for-byte.
                unsafe {
                    leaked.get(concat!($name, "\0").as_bytes()).map_err(|source| {
                        Error::SymbolNotFound { name: $name, source }
                    })?
                }
            }};
        }

        Ok(Arc::new(TfLiteLibrary {
            _lib: lib,

            model_create_from_file: sym!("TfLiteModelCreateFromFile"),
            model_delete: sym!("TfLiteModelDelete"),

            options_create: sym!("TfLiteInterpreterOptionsCreate"),
            options_delete: sym!("TfLiteInterpreterOptionsDelete"),
            options_set_num_threads: sym!("TfLiteInterpreterOptionsSetNumThreads"),
            options_add_delegate: sym!("TfLiteInterpreterOptionsAddDelegate"),

            interpreter_create: sym!("TfLiteInterpreterCreate"),
            interpreter_delete: sym!("TfLiteInterpreterDelete"),
            interpreter_allocate_tensors: sym!("TfLiteInterpreterAllocateTensors"),
            interpreter_invoke: sym!("TfLiteInterpreterInvoke"),
            interpreter_get_input_tensor_count: sym!("TfLiteInterpreterGetInputTensorCount"),
            interpreter_get_output_tensor_count: sym!("TfLiteInterpreterGetOutputTensorCount"),
            interpreter_get_input_tensor: sym!("TfLiteInterpreterGetInputTensor"),
            interpreter_get_output_tensor: sym!("TfLiteInterpreterGetOutputTensor"),

            tensor_type: sym!("TfLiteTensorType"),
            tensor_num_dims: sym!("TfLiteTensorNumDims"),
            tensor_dim: sym!("TfLiteTensorDim"),
            tensor_byte_size: sym!("TfLiteTensorByteSize"),
            tensor_data: sym!("TfLiteTensorData"),
            tensor_name: sym!("TfLiteTensorName"),
            tensor_quantization_params: sym!("TfLiteTensorQuantizationParams"),

            external_delegate: Self::resolve_external_delegate(leaked),
        }))
    }

    // Try to resolve the six experimental external-delegate symbols. Returns
    // `Some` only when every symbol is present; otherwise `None`, and the
    // delegate API surfaces a clean `ExternalDelegateApiUnavailable` error at
    // use time.
    fn resolve_external_delegate(
        leaked: &'static Arc<Library>,
    ) -> Option<ExternalDelegateSymbols> {
        // SAFETY: same justification as the required-symbol resolutions above.
        unsafe {
            let options_create = leaked
                .get::<ExternalDelegateOptionsCreateFn>(b"TfLiteExternalDelegateOptionsCreate\0")
                .ok()?;
            let options_delete = leaked
                .get::<ExternalDelegateOptionsDeleteFn>(b"TfLiteExternalDelegateOptionsDelete\0")
                .ok()?;
            let options_set_library_path = leaked
                .get::<ExternalDelegateOptionsSetLibraryPathFn>(
                    b"TfLiteExternalDelegateOptionsSetLibraryPath\0",
                )
                .ok()?;
            let options_insert = leaked
                .get::<ExternalDelegateOptionsInsertFn>(b"TfLiteExternalDelegateOptionsInsert\0")
                .ok()?;
            let create = leaked
                .get::<ExternalDelegateCreateFn>(b"TfLiteExternalDelegateCreate\0")
                .ok()?;
            let delete = leaked
                .get::<ExternalDelegateDeleteFn>(b"TfLiteExternalDelegateDelete\0")
                .ok()?;
            Some(ExternalDelegateSymbols {
                options_create,
                options_delete,
                options_set_library_path,
                options_insert,
                create,
                delete,
            })
        }
    }
}

/// A delegate built via TFLite's generic *external delegate* mechanism.
///
/// This is how you plug in a vendor-specific accelerator (NNAPI, GPU,
/// XNNPACK, VX/NPU, etc.) without baking knowledge of that accelerator
/// into this crate. You point the builder at the accelerator's plugin
/// shared object (e.g. `libvx_delegate.so`, `libnnapi_delegate.so`,
/// `libxnnpack_delegate.so`), optionally set string key/value options, then
/// call [`ExternalDelegateBuilder::build`] to materialize a delegate handle.
///
/// Pass the resulting [`ExternalDelegate`] to
/// [`crate::interpreter::InterpreterOptions::add_delegate`]; ownership is
/// transferred to TFLite at that point.
pub struct ExternalDelegate {
    pub(crate) raw: *mut TfLiteDelegate,
    pub(crate) lib: Arc<TfLiteLibrary>,
}

// SAFETY: `ExternalDelegate` only owns a raw C handle. It is moved (not
// shared) between threads, and once handed to `InterpreterOptions` it is
// `mem::forget`-ed, so cross-thread aliasing concerns do not arise.
unsafe impl Send for ExternalDelegate {}
unsafe impl Sync for ExternalDelegate {}

impl ExternalDelegate {
    /// Start configuring an external delegate backed by the plugin at `library_path`.
    pub fn builder<P: AsRef<Path>>(
        lib: Arc<TfLiteLibrary>,
        library_path: P,
    ) -> Result<ExternalDelegateBuilder> {
        let path_str = library_path
            .as_ref()
            .to_str()
            .ok_or_else(|| Error::InvalidPath(c_string_nul_error()))?;
        let c_path = CString::new(path_str).map_err(Error::InvalidPath)?;
        Ok(ExternalDelegateBuilder {
            lib,
            library_path: c_path,
            options: Vec::new(),
        })
    }

    /// Returns the raw C handle. Mostly useful for advanced FFI integration.
    pub fn as_ptr(&self) -> *mut TfLiteDelegate {
        self.raw
    }
}

impl Drop for ExternalDelegate {
    fn drop(&mut self) {
        // SAFETY: `raw` was returned by `TfLiteExternalDelegateCreate` and
        // has not yet been handed off to an interpreter (otherwise the
        // owning `InterpreterOptions::add_delegate` would have forgotten
        // us). The external-delegate symbol group must be present — we
        // could only build this struct via `ExternalDelegateBuilder::build`,
        // which already verified that.
        if let Some(ext) = self.lib.external_delegate.as_ref() {
            unsafe {
                (ext.delete)(self.raw);
            }
        }
    }
}

/// Builder for an [`ExternalDelegate`].
pub struct ExternalDelegateBuilder {
    lib: Arc<TfLiteLibrary>,
    library_path: CString,
    options: Vec<(CString, CString)>,
}

impl ExternalDelegateBuilder {
    /// Set a key/value option understood by the underlying delegate plugin.
    ///
    /// Keys and values are UTF-8 strings forwarded verbatim via
    /// `TfLiteExternalDelegateOptionsInsert`. Embedded NUL bytes in either
    /// argument return [`Error::InvalidOption`].
    pub fn with_option(mut self, key: &str, value: &str) -> Result<Self> {
        let k = CString::new(key).map_err(Error::InvalidOption)?;
        let v = CString::new(value).map_err(Error::InvalidOption)?;
        self.options.push((k, v));
        Ok(self)
    }

    /// Materialize the delegate by calling into the TFLite C API.
    ///
    /// Returns [`Error::ExternalDelegateApiUnavailable`] if the loaded
    /// TFLite library does not export the `TfLiteExternalDelegate*` symbol
    /// group (common with stripped or community-built `.so`/`.dylib`).
    pub fn build(self) -> Result<ExternalDelegate> {
        let lib = self.lib;
        let ext = lib
            .external_delegate
            .as_ref()
            .ok_or(Error::ExternalDelegateApiUnavailable)?;

        // SAFETY: Each FFI call below uses the cached symbol with its
        // declared signature. The `options` handle returned by Create is
        // unconditionally Delete-d once we have either a delegate or an
        // error, so it cannot leak.
        unsafe {
            let opts = (ext.options_create)();
            if opts.is_null() {
                return Err(Error::DelegateCreateFailed);
            }

            let set_status = (ext.options_set_library_path)(opts, self.library_path.as_ptr());
            if set_status != TfLiteStatus::Ok {
                (ext.options_delete)(opts);
                return Err(Error::DelegateOptionRejected { status: set_status });
            }

            for (k, v) in &self.options {
                let status = (ext.options_insert)(opts, k.as_ptr(), v.as_ptr());
                if status != TfLiteStatus::Ok {
                    (ext.options_delete)(opts);
                    return Err(Error::DelegateOptionRejected { status });
                }
            }

            let raw = (ext.create)(opts);
            (ext.options_delete)(opts);

            if raw.is_null() {
                return Err(Error::DelegateCreateFailed);
            }
            Ok(ExternalDelegate { raw, lib: lib.clone() })
        }
    }
}

// Internal helper: produce a `NulError` for a path that wasn't valid UTF-8.
// We piggy-back on `CString::new`'s error type because there's no separate
// "non-UTF-8" variant in the crate-public `Error` enum, and embedded NULs
// and non-UTF-8 paths are both fundamentally "this string can't cross the
// C boundary".
fn c_string_nul_error() -> std::ffi::NulError {
    // A single NUL byte is the simplest input that produces a `NulError`.
    CString::new(vec![0u8]).unwrap_err()
}

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

    #[test]
    fn default_paths_are_nonempty_and_named() {
        // The const is statically populated; we just want to confirm the
        // entries we documented are still in the list.
        let len = DEFAULT_LIBRARY_PATHS.len();
        assert!(len > 0, "default path list collapsed to empty");
        // Must contain both the canonical Linux names so users on stock
        // distros and on tensorflow.org builds both have a chance.
        assert!(DEFAULT_LIBRARY_PATHS
            .iter()
            .any(|p| p.ends_with("libtensorflow-lite.so")));
        assert!(DEFAULT_LIBRARY_PATHS
            .iter()
            .any(|p| p.ends_with("libtensorflowlite_c.so")));
    }

    #[test]
    fn load_failed_error_displays_paths() {
        // Construct a synthetic `LoadFailed` so we can exercise `Display`
        // without depending on a real TFLite install.
        let err = Error::LoadFailed {
            source: libloading::Error::DlOpenUnknown,
            paths: vec!["/nope/libtensorflow-lite.so".to_string()],
        };
        let rendered = format!("{err}");
        assert!(rendered.contains("/nope/libtensorflow-lite.so"));
    }

    #[test]
    fn symbol_not_found_error_includes_symbol_name() {
        let err = Error::SymbolNotFound {
            name: "TfLiteSomeFunction",
            source: libloading::Error::DlSymUnknown,
        };
        assert!(format!("{err}").contains("TfLiteSomeFunction"));
    }
}