Skip to main content

tflite_c/
model.rs

1//! Loaded `TfLiteModel` handle.
2
3use std::ffi::CString;
4use std::path::Path;
5use std::sync::Arc;
6
7use crate::error::{Error, Result};
8use crate::ffi::TfLiteModel;
9use crate::library::TfLiteLibrary;
10
11/// A TensorFlow Lite model loaded from a `.tflite` flatbuffer file.
12///
13/// The underlying `TfLiteModel*` is released when this value is dropped.
14/// Cloning the [`Arc<TfLiteLibrary>`] is cheap and keeps the loaded shared
15/// object alive for as long as any [`Model`] or interpreter still uses it.
16pub struct Model {
17    raw: *mut TfLiteModel,
18    lib: Arc<TfLiteLibrary>,
19}
20
21// SAFETY: A `TfLiteModel` is read-only after construction and the upstream
22// C API is documented as safe to share across threads. The wrapper itself
23// is just a raw pointer plus an `Arc`, both of which we mark explicitly.
24unsafe impl Send for Model {}
25unsafe impl Sync for Model {}
26
27impl Model {
28    /// Load a `.tflite` flatbuffer from disk.
29    ///
30    /// Internally calls `TfLiteModelCreateFromFile`. Returns
31    /// [`Error::ModelLoadFailed`] if the C API returns NULL (commonly: the
32    /// file does not exist, is unreadable, or is not a valid TFLite model).
33    pub fn from_file<P: AsRef<Path>>(path: P, lib: Arc<TfLiteLibrary>) -> Result<Self> {
34        let path_ref = path.as_ref();
35        let path_str = path_ref
36            .to_str()
37            .ok_or_else(|| Error::ModelLoadFailed { path: path_ref.to_path_buf() })?;
38        let c_path = CString::new(path_str).map_err(Error::InvalidPath)?;
39
40        // SAFETY: `c_path` lives for the duration of this call; the C API
41        // copies the string contents it needs.
42        let raw = unsafe { (lib.model_create_from_file)(c_path.as_ptr()) };
43        if raw.is_null() {
44            return Err(Error::ModelLoadFailed { path: path_ref.to_path_buf() });
45        }
46        Ok(Self { raw, lib })
47    }
48
49    /// Returns the raw `TfLiteModel*`. Mostly useful for advanced FFI scenarios.
50    pub fn as_ptr(&self) -> *const TfLiteModel {
51        self.raw
52    }
53
54    pub(crate) fn lib(&self) -> &Arc<TfLiteLibrary> {
55        &self.lib
56    }
57}
58
59impl Drop for Model {
60    fn drop(&mut self) {
61        // SAFETY: `raw` was allocated by `TfLiteModelCreateFromFile` and is
62        // not aliased elsewhere — the interpreter borrows but does not take
63        // ownership.
64        unsafe {
65            (self.lib.model_delete)(self.raw);
66        }
67    }
68}