edgefirst_tflite/

delegate.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2025 Au-Zone Technologies. All Rights Reserved.

//! Delegate loading with configuration options.
//!
//! Delegates provide hardware acceleration for `TFLite` inference. The most
//! common delegate for i.MX platforms is the `VxDelegate`, which offloads
//! operations to the NPU.

use std::ffi::{c_void, CString};
use std::path::Path;
use std::ptr::{self, NonNull};

use edgefirst_tflite_sys::TfLiteDelegate;

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

#[cfg(feature = "dmabuf")]
use edgefirst_tflite_sys::hal_ffi::HalDmaBufFunctions;

#[cfg(feature = "dmabuf")]
use edgefirst_tflite_sys::vx_ffi::VxDmaBufFunctions;

#[cfg(feature = "camera_adaptor")]
use edgefirst_tflite_sys::hal_ffi::HalCameraAdaptorFunctions;

#[cfg(feature = "camera_adaptor")]
use edgefirst_tflite_sys::vx_ffi::VxCameraAdaptorFunctions;

// ---------------------------------------------------------------------------
// DelegateOptions
// ---------------------------------------------------------------------------

/// Key-value options for configuring an external delegate.
///
/// # Examples
///
/// ```no_run
/// use edgefirst_tflite::DelegateOptions;
///
/// let opts = DelegateOptions::new()
///     .option("cache_file_path", "/tmp/vx_cache")
///     .option("device_id", "0");
/// ```
#[derive(Debug, Default, Clone)]
pub struct DelegateOptions {
    options: Vec<(String, String)>,
}

impl DelegateOptions {
    /// Create an empty set of delegate options.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a key-value option pair.
    #[must_use]
    pub fn option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.options.push((key.into(), value.into()));
        self
    }
}

// ---------------------------------------------------------------------------
// Delegate
// ---------------------------------------------------------------------------

/// An external `TFLite` delegate for hardware acceleration.
///
/// Delegates are loaded from shared libraries that export the standard
/// `tflite_plugin_create_delegate` / `tflite_plugin_destroy_delegate`
/// entry points.
///
/// # Examples
///
/// ```no_run
/// use edgefirst_tflite::{Delegate, DelegateOptions};
///
/// // Load delegate with default options
/// let delegate = Delegate::load("libvx_delegate.so")?;
///
/// // Load delegate with options
/// let delegate = Delegate::load_with_options(
///     "libvx_delegate.so",
///     &DelegateOptions::new()
///         .option("cache_file_path", "/tmp/vx_cache")
///         .option("device_id", "0"),
/// )?;
/// # Ok::<(), edgefirst_tflite::Error>(())
/// ```
#[allow(clippy::struct_field_names)]
pub struct Delegate {
    delegate: NonNull<TfLiteDelegate>,
    free: unsafe extern "C" fn(*mut TfLiteDelegate),
    // Keeps the delegate .so loaded for the delegate's lifetime.
    _lib: libloading::Library,

    #[cfg(feature = "dmabuf")]
    hal_dmabuf_fns: Option<HalDmaBufFunctions>,

    /// Inner delegate handle returned by `hal_dmabuf_get_instance()`.
    ///
    /// This is the opaque `hal_delegate_t` (`*mut c_void`) that HAL API
    /// functions expect as their first argument. It is distinct from the
    /// `TfLiteDelegate*` outer pointer and must be used for all HAL calls.
    /// Both `DmaBuf` and `CameraAdaptor` share this same handle.
    #[cfg(feature = "dmabuf")]
    hal_delegate_handle: Option<*mut c_void>,

    #[cfg(feature = "dmabuf")]
    dmabuf_fns: Option<VxDmaBufFunctions>,

    #[cfg(feature = "camera_adaptor")]
    hal_camera_fns: Option<HalCameraAdaptorFunctions>,

    #[cfg(feature = "camera_adaptor")]
    camera_adaptor_fns: Option<VxCameraAdaptorFunctions>,
}

// SAFETY: `hal_delegate_handle` is a raw pointer obtained from
// `hal_dmabuf_get_instance()`. The HAL contract guarantees this pointer
// is valid and stable for the lifetime of the loaded delegate library.
// `Delegate` is the sole owner and never shares the handle concurrently.
#[cfg(feature = "dmabuf")]
// SAFETY: See above — the handle is stable and not concurrently accessed.
unsafe impl Send for Delegate {}
#[cfg(feature = "dmabuf")]
// SAFETY: All HAL API methods take `&self` and the underlying C functions
// are thread-safe per the HAL contract.
unsafe impl Sync for Delegate {}

impl Delegate {
    /// Load an external delegate from a shared library with default options.
    ///
    /// # Errors
    ///
    /// Returns an error if the library cannot be loaded, required symbols
    /// are missing, or the delegate returns a null pointer.
    pub fn load(path: impl AsRef<Path>) -> Result<Self> {
        Self::load_with_options(path, &DelegateOptions::default())
    }

    /// Load an external delegate with configuration options.
    ///
    /// # Errors
    ///
    /// Returns an error if the library cannot be loaded, required symbols
    /// are missing, the delegate returns a null pointer, or any option key
    /// or value contains an interior NUL byte.
    pub fn load_with_options(path: impl AsRef<Path>, options: &DelegateOptions) -> Result<Self> {
        // SAFETY: Loading the shared library via `libloading`. The library is
        // kept alive in `_lib` for the lifetime of the `Delegate`.
        let lib =
            unsafe { libloading::Library::new(path.as_ref().as_os_str()) }.map_err(Error::from)?;

        // SAFETY: Resolving the `tflite_plugin_create_delegate` symbol from
        // the loaded library. The library is valid and loaded above.
        let create_fn = unsafe {
            lib.get::<unsafe extern "C" fn(
                *const *const std::os::raw::c_char,
                *const *const std::os::raw::c_char,
                usize,
                Option<unsafe extern "C" fn(*const std::os::raw::c_char)>,
            ) -> *mut TfLiteDelegate>(b"tflite_plugin_create_delegate")
        }
        .map_err(Error::from)?;

        // SAFETY: Resolving the `tflite_plugin_destroy_delegate` symbol from
        // the same loaded library.
        let destroy_fn = unsafe {
            lib.get::<unsafe extern "C" fn(*mut TfLiteDelegate)>(b"tflite_plugin_destroy_delegate")
        }
        .map_err(Error::from)?;

        // Convert options to C string arrays.
        let (keys_c, values_c): (Vec<CString>, Vec<CString>) = options
            .options
            .iter()
            .map(|(k, v)| {
                Ok((
                    CString::new(k.as_str()).map_err(|_| {
                        Error::invalid_argument(format!(
                            "option key \"{k}\" contains interior NUL byte"
                        ))
                    })?,
                    CString::new(v.as_str()).map_err(|_| {
                        Error::invalid_argument(format!(
                            "option value \"{v}\" contains interior NUL byte"
                        ))
                    })?,
                ))
            })
            .collect::<Result<Vec<_>>>()?
            .into_iter()
            .unzip();
        let keys_ptrs: Vec<*const std::os::raw::c_char> =
            keys_c.iter().map(|c| c.as_ptr()).collect();
        let values_ptrs: Vec<*const std::os::raw::c_char> =
            values_c.iter().map(|c| c.as_ptr()).collect();

        // SAFETY: `create_fn` is a valid symbol resolved above. `keys_ptrs`
        // and `values_ptrs` point to valid NUL-terminated C strings (from
        // `CString::new`), or null when empty. `keys_c` and `values_c` are
        // alive for this call, keeping the pointers valid.
        let raw = unsafe {
            create_fn(
                if keys_ptrs.is_empty() {
                    ptr::null()
                } else {
                    keys_ptrs.as_ptr()
                },
                if values_ptrs.is_empty() {
                    ptr::null()
                } else {
                    values_ptrs.as_ptr()
                },
                options.options.len(),
                None,
            )
        };

        let delegate = NonNull::new(raw)
            .ok_or_else(|| Error::null_pointer("tflite_plugin_create_delegate returned null"))?;

        // Copy the destroy function pointer before lib is stored.
        let free = *destroy_fn;

        // Probe for delegate DMA-BUF extensions.
        #[cfg(feature = "dmabuf")]
        // SAFETY: `lib` is a valid loaded library. `try_load` resolves
        // optional symbols; missing symbols return `None`, not UB.
        let hal_dmabuf_fns = unsafe { HalDmaBufFunctions::try_load(&lib) };

        // Call `hal_dmabuf_get_instance()` to obtain the inner delegate handle.
        // This is the opaque `hal_delegate_t` pointer that all HAL API calls
        // expect. A null result means HAL is not available on this device.
        #[cfg(feature = "dmabuf")]
        let hal_delegate_handle: Option<*mut c_void> = hal_dmabuf_fns.as_ref().and_then(|fns| {
            // SAFETY: `get_instance` is a valid function pointer loaded from
            // the delegate library. It takes no arguments and returns an opaque
            // handle that is valid for the lifetime of the library.
            let ptr = unsafe { (fns.get_instance)() };
            if ptr.is_null() {
                None
            } else {
                Some(ptr)
            }
        });

        #[cfg(feature = "dmabuf")]
        // SAFETY: Same as above — resolves optional VxDelegate DMA-BUF
        // symbols as a fallback for delegates that haven't adopted the
        // HAL DMA-BUF API yet.
        let dmabuf_fns = unsafe { VxDmaBufFunctions::try_load(&lib) };

        #[cfg(feature = "camera_adaptor")]
        // SAFETY: Same as above — resolves optional HAL Camera Adaptor
        // symbols from the loaded library.
        let hal_camera_fns = unsafe { HalCameraAdaptorFunctions::try_load(&lib) };

        #[cfg(feature = "camera_adaptor")]
        // SAFETY: Same as `VxDmaBufFunctions::try_load` above — resolves
        // optional CameraAdaptor symbols from the loaded library.
        let camera_adaptor_fns = unsafe { VxCameraAdaptorFunctions::try_load(&lib) };

        Ok(Self {
            delegate,
            free,
            _lib: lib,
            #[cfg(feature = "dmabuf")]
            hal_dmabuf_fns,
            #[cfg(feature = "dmabuf")]
            hal_delegate_handle,
            #[cfg(feature = "dmabuf")]
            dmabuf_fns,
            #[cfg(feature = "camera_adaptor")]
            hal_camera_fns,
            #[cfg(feature = "camera_adaptor")]
            camera_adaptor_fns,
        })
    }

    /// Returns the raw delegate pointer.
    ///
    /// This is an escape hatch for advanced use cases that need direct
    /// FFI access to the delegate.
    #[must_use]
    pub fn as_ptr(&self) -> *mut TfLiteDelegate {
        self.delegate.as_ptr()
    }

    /// Access DMA-BUF extensions if available on this delegate.
    ///
    /// Returns `Some` if the delegate exports either the HAL Delegate
    /// DMA-BUF API (`hal_dmabuf_*`) or the legacy `VxDelegate` DMA-BUF API.
    /// The HAL API is preferred when both are available.
    #[cfg(feature = "dmabuf")]
    #[must_use]
    pub fn dmabuf(&self) -> Option<crate::dmabuf::DmaBuf<'_>> {
        if self.hal_dmabuf_fns.is_some() || self.dmabuf_fns.is_some() {
            Some(crate::dmabuf::DmaBuf::new(
                self.delegate,
                self.hal_delegate_handle,
                self.hal_dmabuf_fns.as_ref(),
                self.dmabuf_fns.as_ref(),
            ))
        } else {
            None
        }
    }

    /// Returns `true` if this delegate supports DMA-BUF zero-copy.
    #[cfg(feature = "dmabuf")]
    #[must_use]
    pub fn has_dmabuf(&self) -> bool {
        self.hal_dmabuf_fns.is_some() || self.dmabuf_fns.is_some()
    }

    /// Access `CameraAdaptor` extensions if available on this delegate.
    ///
    /// Returns `Some` if the delegate exports either the HAL Delegate
    /// Camera Adaptor API (`hal_camera_adaptor_*`) or the legacy
    /// `VxDelegate` `CameraAdaptor` API. The HAL API is preferred when
    /// both are available.
    #[cfg(feature = "camera_adaptor")]
    #[must_use]
    pub fn camera_adaptor(&self) -> Option<crate::camera_adaptor::CameraAdaptor<'_>> {
        if self.hal_camera_fns.is_some() || self.camera_adaptor_fns.is_some() {
            // The CameraAdaptor HAL API reuses the same inner delegate handle
            // as the DMA-BUF HAL API — there is no separate
            // `hal_camera_adaptor_get_instance`.
            #[cfg(feature = "dmabuf")]
            let hal_handle = self.hal_delegate_handle;
            #[cfg(not(feature = "dmabuf"))]
            let hal_handle: Option<*mut std::ffi::c_void> = None;

            Some(crate::camera_adaptor::CameraAdaptor::new(
                self.delegate,
                hal_handle,
                self.hal_camera_fns.as_ref(),
                self.camera_adaptor_fns.as_ref(),
            ))
        } else {
            None
        }
    }

    /// Returns `true` if this delegate supports `CameraAdaptor`.
    #[cfg(feature = "camera_adaptor")]
    #[must_use]
    pub fn has_camera_adaptor(&self) -> bool {
        self.hal_camera_fns.is_some() || self.camera_adaptor_fns.is_some()
    }
}

#[allow(clippy::missing_fields_in_debug)]
impl std::fmt::Debug for Delegate {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut d = f.debug_struct("Delegate");
        d.field("ptr", &self.delegate);

        #[cfg(feature = "dmabuf")]
        d.field("has_hal_dmabuf", &self.hal_dmabuf_fns.is_some());

        #[cfg(feature = "dmabuf")]
        d.field(
            "has_hal_delegate_handle",
            &self.hal_delegate_handle.is_some(),
        );

        #[cfg(feature = "dmabuf")]
        d.field("has_vx_dmabuf", &self.dmabuf_fns.is_some());

        #[cfg(feature = "camera_adaptor")]
        d.field("has_hal_camera_adaptor", &self.hal_camera_fns.is_some());

        #[cfg(feature = "camera_adaptor")]
        d.field("has_vx_camera_adaptor", &self.camera_adaptor_fns.is_some());

        d.finish_non_exhaustive()
    }
}

impl Drop for Delegate {
    fn drop(&mut self) {
        // SAFETY: The delegate pointer was created by `tflite_plugin_create_delegate`
        // and `free` is the matching `tflite_plugin_destroy_delegate` from the same
        // library, which is still loaded (held by `_lib`).
        unsafe { (self.free)(self.delegate.as_ptr()) };
    }
}

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

    #[test]
    fn new_creates_empty_options() {
        let opts = DelegateOptions::new();
        let debug = format!("{opts:?}");
        assert_eq!(debug, "DelegateOptions { options: [] }");
    }

    #[test]
    fn builder_chaining() {
        let opts = DelegateOptions::new().option("a", "1").option("b", "2");
        assert_eq!(opts.options.len(), 2);
    }

    #[test]
    fn default_matches_new() {
        let from_new = format!("{:?}", DelegateOptions::new());
        let from_default = format!("{:?}", DelegateOptions::default());
        assert_eq!(from_new, from_default);
    }

    #[test]
    fn clone_produces_equal_values() {
        let opts = DelegateOptions::new().option("key", "value");
        let cloned = opts.clone();
        assert_eq!(format!("{opts:?}"), format!("{cloned:?}"));
    }

    #[test]
    fn debug_formatting_not_empty() {
        let opts = DelegateOptions::new().option("cache", "/tmp");
        let debug = format!("{opts:?}");
        assert!(!debug.is_empty());
        assert!(debug.contains("DelegateOptions"));
        assert!(debug.contains("cache"));
        assert!(debug.contains("/tmp"));
    }
}