Skip to main content

llama_cpp_bindings/mtmd/
mtmd_bitmap.rs

1use std::ffi::{CStr, CString, c_char};
2use std::path::PathBuf;
3use std::ptr::NonNull;
4use std::slice;
5
6use crate::ffi_error_reader::read_and_free_cpp_error;
7
8use super::mtmd_bitmap_error::MtmdBitmapError;
9use super::mtmd_context::MtmdContext;
10
11fn cstr_ptr_to_optional_string(ptr: *const c_char) -> Option<String> {
12    if ptr.is_null() {
13        None
14    } else {
15        let id = unsafe { CStr::from_ptr(ptr) }
16            .to_string_lossy()
17            .into_owned();
18
19        Some(id)
20    }
21}
22
23/// # Safety
24///
25/// `out_bitmap` must be either null or a valid pointer to an `mtmd_bitmap`
26/// allocated by `llama_rs_mtmd_bitmap_init_from_file`. `out_error` must be
27/// either null or a valid pointer to a null-terminated C string allocated by
28/// `llama_rs_dup_string`.
29unsafe fn from_file_status_to_result(
30    status: llama_cpp_bindings_sys::llama_rs_mtmd_bitmap_init_from_file_status,
31    out_bitmap: *mut llama_cpp_bindings_sys::mtmd_bitmap,
32    out_error: *mut c_char,
33    path: &str,
34) -> Result<MtmdBitmap, MtmdBitmapError> {
35    match status {
36        llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_OK => {
37            let bitmap = NonNull::new(out_bitmap).ok_or_else(|| {
38                MtmdBitmapError::FileUnreadable {
39                    path: PathBuf::from(path),
40                }
41            })?;
42            Ok(MtmdBitmap { bitmap })
43        }
44        llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_RETURNED_NULL => {
45            Err(MtmdBitmapError::FileUnreadable {
46                path: PathBuf::from(path),
47            })
48        }
49        llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_ERROR_STRING_ALLOCATION_FAILED => {
50            Err(MtmdBitmapError::NotEnoughMemory)
51        }
52        llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_THREW_CXX_EXCEPTION => {
53            let message = unsafe { read_and_free_cpp_error(out_error) };
54            Err(MtmdBitmapError::Reported { message })
55        }
56        other => unreachable!(
57            "llama_rs_mtmd_bitmap_init_from_file returned unrecognized status: {other}"
58        ),
59    }
60}
61
62#[derive(Debug, Clone)]
63pub struct MtmdBitmap {
64    pub bitmap: NonNull<llama_cpp_bindings_sys::mtmd_bitmap>,
65}
66
67unsafe impl Send for MtmdBitmap {}
68unsafe impl Sync for MtmdBitmap {}
69
70impl MtmdBitmap {
71    /// # Errors
72    ///
73    /// * `InvalidDataSize` - Data length doesn't match `nx * ny * 3`
74    /// * `NullResult` - Underlying C function returned null
75    ///
76    pub fn from_image_data(nx: u32, ny: u32, data: &[u8]) -> Result<Self, MtmdBitmapError> {
77        if nx < 2 || ny < 2 {
78            return Err(MtmdBitmapError::ImageDimensionsTooSmall(nx, ny));
79        }
80
81        if data.len() != (nx * ny * 3) as usize {
82            return Err(MtmdBitmapError::InvalidDataSize);
83        }
84
85        let bitmap = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_init(nx, ny, data.as_ptr()) };
86
87        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::BitmapDecodeFailed)?;
88
89        Ok(Self { bitmap })
90    }
91
92    /// # Errors
93    ///
94    /// * `NullResult` - Underlying C function returned null
95    ///
96    pub fn from_audio_data(data: &[f32]) -> Result<Self, MtmdBitmapError> {
97        let bitmap = unsafe {
98            llama_cpp_bindings_sys::mtmd_bitmap_init_from_audio(data.len(), data.as_ptr())
99        };
100
101        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::BitmapDecodeFailed)?;
102
103        Ok(Self { bitmap })
104    }
105
106    /// # Errors
107    ///
108    /// Returns an [`MtmdBitmapError`] variant matching the wrapper's status code.
109    pub fn from_file(ctx: &MtmdContext, path: &str) -> Result<Self, MtmdBitmapError> {
110        let path_cstr = CString::new(path)?;
111        let mut out_bitmap: *mut llama_cpp_bindings_sys::mtmd_bitmap = std::ptr::null_mut();
112        let mut out_error: *mut c_char = std::ptr::null_mut();
113
114        let status = unsafe {
115            llama_cpp_bindings_sys::llama_rs_mtmd_bitmap_init_from_file(
116                ctx.context.as_ptr(),
117                path_cstr.as_ptr(),
118                &raw mut out_bitmap,
119                &raw mut out_error,
120            )
121        };
122
123        unsafe { from_file_status_to_result(status, out_bitmap, out_error, path) }
124    }
125
126    /// # Errors
127    ///
128    /// * `NullResult` - Buffer could not be processed
129    pub fn from_buffer(ctx: &MtmdContext, data: &[u8]) -> Result<Self, MtmdBitmapError> {
130        let bitmap_wrapper = unsafe {
131            llama_cpp_bindings_sys::mtmd_helper_bitmap_init_from_buf(
132                ctx.context.as_ptr(),
133                data.as_ptr(),
134                data.len(),
135                false,
136            )
137        };
138
139        let bitmap =
140            NonNull::new(bitmap_wrapper.bitmap).ok_or(MtmdBitmapError::BitmapDecodeFailed)?;
141
142        Ok(Self { bitmap })
143    }
144
145    #[must_use]
146    pub fn nx(&self) -> u32 {
147        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_nx(self.bitmap.as_ptr()) }
148    }
149
150    #[must_use]
151    pub fn ny(&self) -> u32 {
152        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_ny(self.bitmap.as_ptr()) }
153    }
154
155    #[must_use]
156    pub fn data(&self) -> &[u8] {
157        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_data(self.bitmap.as_ptr()) };
158        let len = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_n_bytes(self.bitmap.as_ptr()) };
159        unsafe { slice::from_raw_parts(ptr, len) }
160    }
161
162    #[must_use]
163    pub fn is_audio(&self) -> bool {
164        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_is_audio(self.bitmap.as_ptr()) }
165    }
166
167    #[must_use]
168    pub fn id(&self) -> Option<String> {
169        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_id(self.bitmap.as_ptr()) };
170
171        cstr_ptr_to_optional_string(ptr)
172    }
173
174    /// # Errors
175    ///
176    /// Returns an error if the ID string contains null bytes.
177    ///
178    pub fn set_id(&self, id: &str) -> Result<(), std::ffi::NulError> {
179        let id_cstr = CString::new(id)?;
180        unsafe {
181            llama_cpp_bindings_sys::mtmd_bitmap_set_id(self.bitmap.as_ptr(), id_cstr.as_ptr());
182        }
183
184        Ok(())
185    }
186}
187
188impl Drop for MtmdBitmap {
189    fn drop(&mut self) {
190        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_free(self.bitmap.as_ptr()) }
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use std::path::PathBuf;
197
198    use super::MtmdBitmap;
199    use super::MtmdBitmapError;
200
201    #[test]
202    fn cstr_ptr_to_optional_string_returns_none_for_null() {
203        assert!(super::cstr_ptr_to_optional_string(std::ptr::null()).is_none());
204    }
205
206    #[test]
207    fn cstr_ptr_to_optional_string_returns_some_for_valid() {
208        let cstr = std::ffi::CString::new("hello").unwrap();
209        let result = super::cstr_ptr_to_optional_string(cstr.as_ptr());
210
211        assert_eq!(result, Some("hello".to_string()));
212    }
213
214    #[test]
215    fn from_image_data_creates_valid_bitmap() {
216        let red_pixel: [u8; 3] = [255, 0, 0];
217        let image_data: Vec<u8> = red_pixel.repeat(4);
218        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
219        assert_eq!(bitmap.nx(), 2);
220        assert_eq!(bitmap.ny(), 2);
221        assert!(!bitmap.is_audio());
222    }
223
224    #[test]
225    fn invalid_data_size_returns_error() {
226        let too_short = vec![0u8; 5];
227        let result = MtmdBitmap::from_image_data(2, 2, &too_short);
228        assert!(result.is_err());
229    }
230
231    #[test]
232    fn from_image_data_rejects_dimensions_below_minimum() {
233        let result_1x1 = MtmdBitmap::from_image_data(1, 1, &[0u8; 3]);
234        let result_1x2 = MtmdBitmap::from_image_data(1, 2, &[0u8; 6]);
235        let result_2x1 = MtmdBitmap::from_image_data(2, 1, &[0u8; 6]);
236        let result_0x0 = MtmdBitmap::from_image_data(0, 0, &[]);
237
238        assert_eq!(
239            result_1x1.unwrap_err(),
240            MtmdBitmapError::ImageDimensionsTooSmall(1, 1)
241        );
242        assert_eq!(
243            result_1x2.unwrap_err(),
244            MtmdBitmapError::ImageDimensionsTooSmall(1, 2)
245        );
246        assert_eq!(
247            result_2x1.unwrap_err(),
248            MtmdBitmapError::ImageDimensionsTooSmall(2, 1)
249        );
250        assert_eq!(
251            result_0x0.unwrap_err(),
252            MtmdBitmapError::ImageDimensionsTooSmall(0, 0)
253        );
254    }
255
256    #[test]
257    fn set_id_changes_id() {
258        let image_data = vec![0u8; 12];
259        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
260        bitmap.set_id("test_image").unwrap();
261
262        assert_eq!(bitmap.id().as_deref(), Some("test_image"));
263    }
264
265    #[test]
266    fn from_audio_data_creates_valid_bitmap() {
267        let audio_samples: Vec<f32> = (0u8..100)
268            .map(|index| (f32::from(index) * 0.1).sin())
269            .collect();
270        let bitmap = MtmdBitmap::from_audio_data(&audio_samples).unwrap();
271
272        assert!(bitmap.is_audio());
273    }
274
275    #[test]
276    fn data_returns_expected_bytes_for_image() {
277        let pixel_data: Vec<u8> = vec![255, 0, 0, 0, 255, 0, 0, 0, 255, 128, 128, 128];
278        let bitmap = MtmdBitmap::from_image_data(2, 2, &pixel_data).unwrap();
279        let returned_data = bitmap.data();
280
281        assert_eq!(returned_data, &pixel_data);
282    }
283
284    #[test]
285    fn id_returns_some_by_default() {
286        let image_data = vec![0u8; 12];
287        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
288
289        assert!(bitmap.id().is_some());
290    }
291
292    #[test]
293    fn id_returns_custom_value_after_set() {
294        let image_data = vec![0u8; 12];
295        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
296        bitmap.set_id("my_image").unwrap();
297
298        assert_eq!(bitmap.id(), Some("my_image".to_string()));
299    }
300
301    #[test]
302    fn set_id_with_null_byte_returns_error() {
303        let image_data = vec![0u8; 12];
304        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
305        let result = bitmap.set_id("id\0null");
306
307        assert!(result.is_err());
308    }
309
310    #[test]
311    fn from_file_status_ok_with_null_bitmap_returns_file_unreadable() {
312        let result = unsafe {
313            super::from_file_status_to_result(
314                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_OK,
315                std::ptr::null_mut(),
316                std::ptr::null_mut(),
317                "/missing/image.png",
318            )
319        };
320
321        assert_eq!(
322            result.unwrap_err(),
323            MtmdBitmapError::FileUnreadable {
324                path: PathBuf::from("/missing/image.png")
325            }
326        );
327    }
328
329    #[test]
330    fn from_file_status_vendored_returned_null_returns_file_unreadable() {
331        let result = unsafe {
332            super::from_file_status_to_result(
333                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_RETURNED_NULL,
334                std::ptr::null_mut(),
335                std::ptr::null_mut(),
336                "/missing/image.png",
337            )
338        };
339
340        assert_eq!(
341            result.unwrap_err(),
342            MtmdBitmapError::FileUnreadable {
343                path: PathBuf::from("/missing/image.png")
344            }
345        );
346    }
347
348    #[test]
349    fn from_file_status_error_string_allocation_failed_returns_not_enough_memory() {
350        let result = unsafe {
351            super::from_file_status_to_result(
352                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_ERROR_STRING_ALLOCATION_FAILED,
353                std::ptr::null_mut(),
354                std::ptr::null_mut(),
355                "/missing/image.png",
356            )
357        };
358
359        assert_eq!(result.unwrap_err(), MtmdBitmapError::NotEnoughMemory);
360    }
361
362    #[test]
363    fn from_file_status_vendored_threw_cxx_exception_returns_reported() {
364        let result = unsafe {
365            super::from_file_status_to_result(
366                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_THREW_CXX_EXCEPTION,
367                std::ptr::null_mut(),
368                std::ptr::null_mut(),
369                "/missing/image.png",
370            )
371        };
372
373        assert_eq!(
374            result.unwrap_err(),
375            MtmdBitmapError::Reported {
376                message: "unknown error".to_string()
377            }
378        );
379    }
380
381    #[test]
382    #[should_panic(expected = "returned unrecognized status")]
383    fn from_file_status_null_ctx_arg_panics_as_unreachable() {
384        let _result = unsafe {
385            super::from_file_status_to_result(
386                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_NULL_CTX_ARG,
387                std::ptr::null_mut(),
388                std::ptr::null_mut(),
389                "/missing/image.png",
390            )
391        };
392    }
393}