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 = 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            )
136        };
137
138        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::BitmapDecodeFailed)?;
139
140        Ok(Self { bitmap })
141    }
142
143    #[must_use]
144    pub fn nx(&self) -> u32 {
145        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_nx(self.bitmap.as_ptr()) }
146    }
147
148    #[must_use]
149    pub fn ny(&self) -> u32 {
150        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_ny(self.bitmap.as_ptr()) }
151    }
152
153    #[must_use]
154    pub fn data(&self) -> &[u8] {
155        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_data(self.bitmap.as_ptr()) };
156        let len = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_n_bytes(self.bitmap.as_ptr()) };
157        unsafe { slice::from_raw_parts(ptr, len) }
158    }
159
160    #[must_use]
161    pub fn is_audio(&self) -> bool {
162        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_is_audio(self.bitmap.as_ptr()) }
163    }
164
165    #[must_use]
166    pub fn id(&self) -> Option<String> {
167        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_id(self.bitmap.as_ptr()) };
168
169        cstr_ptr_to_optional_string(ptr)
170    }
171
172    /// # Errors
173    ///
174    /// Returns an error if the ID string contains null bytes.
175    ///
176    pub fn set_id(&self, id: &str) -> Result<(), std::ffi::NulError> {
177        let id_cstr = CString::new(id)?;
178        unsafe {
179            llama_cpp_bindings_sys::mtmd_bitmap_set_id(self.bitmap.as_ptr(), id_cstr.as_ptr());
180        }
181
182        Ok(())
183    }
184}
185
186impl Drop for MtmdBitmap {
187    fn drop(&mut self) {
188        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_free(self.bitmap.as_ptr()) }
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use std::path::PathBuf;
195
196    use super::MtmdBitmap;
197    use super::MtmdBitmapError;
198
199    #[test]
200    fn cstr_ptr_to_optional_string_returns_none_for_null() {
201        assert!(super::cstr_ptr_to_optional_string(std::ptr::null()).is_none());
202    }
203
204    #[test]
205    fn cstr_ptr_to_optional_string_returns_some_for_valid() {
206        let cstr = std::ffi::CString::new("hello").unwrap();
207        let result = super::cstr_ptr_to_optional_string(cstr.as_ptr());
208
209        assert_eq!(result, Some("hello".to_string()));
210    }
211
212    #[test]
213    fn from_image_data_creates_valid_bitmap() {
214        let red_pixel: [u8; 3] = [255, 0, 0];
215        let image_data: Vec<u8> = red_pixel.repeat(4);
216        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
217        assert_eq!(bitmap.nx(), 2);
218        assert_eq!(bitmap.ny(), 2);
219        assert!(!bitmap.is_audio());
220    }
221
222    #[test]
223    fn invalid_data_size_returns_error() {
224        let too_short = vec![0u8; 5];
225        let result = MtmdBitmap::from_image_data(2, 2, &too_short);
226        assert!(result.is_err());
227    }
228
229    #[test]
230    fn from_image_data_rejects_dimensions_below_minimum() {
231        let result_1x1 = MtmdBitmap::from_image_data(1, 1, &[0u8; 3]);
232        let result_1x2 = MtmdBitmap::from_image_data(1, 2, &[0u8; 6]);
233        let result_2x1 = MtmdBitmap::from_image_data(2, 1, &[0u8; 6]);
234        let result_0x0 = MtmdBitmap::from_image_data(0, 0, &[]);
235
236        assert_eq!(
237            result_1x1.unwrap_err(),
238            MtmdBitmapError::ImageDimensionsTooSmall(1, 1)
239        );
240        assert_eq!(
241            result_1x2.unwrap_err(),
242            MtmdBitmapError::ImageDimensionsTooSmall(1, 2)
243        );
244        assert_eq!(
245            result_2x1.unwrap_err(),
246            MtmdBitmapError::ImageDimensionsTooSmall(2, 1)
247        );
248        assert_eq!(
249            result_0x0.unwrap_err(),
250            MtmdBitmapError::ImageDimensionsTooSmall(0, 0)
251        );
252    }
253
254    #[test]
255    fn set_id_changes_id() {
256        let image_data = vec![0u8; 12];
257        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
258        bitmap.set_id("test_image").unwrap();
259
260        assert_eq!(bitmap.id().as_deref(), Some("test_image"));
261    }
262
263    #[test]
264    fn from_audio_data_creates_valid_bitmap() {
265        #[expect(
266            clippy::cast_precision_loss,
267            reason = "test fixture casts a small i32 (0..100) to f32 to synthesise a sine wave; \
268                      the values are well within f32's exact-representation range"
269        )]
270        let audio_samples: Vec<f32> = (0..100).map(|index| (index as f32 * 0.1).sin()).collect();
271        let bitmap = MtmdBitmap::from_audio_data(&audio_samples).unwrap();
272
273        assert!(bitmap.is_audio());
274    }
275
276    #[test]
277    fn data_returns_expected_bytes_for_image() {
278        let pixel_data: Vec<u8> = vec![255, 0, 0, 0, 255, 0, 0, 0, 255, 128, 128, 128];
279        let bitmap = MtmdBitmap::from_image_data(2, 2, &pixel_data).unwrap();
280        let returned_data = bitmap.data();
281
282        assert_eq!(returned_data, &pixel_data);
283    }
284
285    #[test]
286    fn id_returns_some_by_default() {
287        let image_data = vec![0u8; 12];
288        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
289
290        assert!(bitmap.id().is_some());
291    }
292
293    #[test]
294    fn id_returns_custom_value_after_set() {
295        let image_data = vec![0u8; 12];
296        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
297        bitmap.set_id("my_image").unwrap();
298
299        assert_eq!(bitmap.id(), Some("my_image".to_string()));
300    }
301
302    #[test]
303    fn set_id_with_null_byte_returns_error() {
304        let image_data = vec![0u8; 12];
305        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
306        let result = bitmap.set_id("id\0null");
307
308        assert!(result.is_err());
309    }
310
311    #[test]
312    fn from_file_status_ok_with_null_bitmap_returns_file_unreadable() {
313        let result = unsafe {
314            super::from_file_status_to_result(
315                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_OK,
316                std::ptr::null_mut(),
317                std::ptr::null_mut(),
318                "/missing/image.png",
319            )
320        };
321
322        assert_eq!(
323            result.unwrap_err(),
324            MtmdBitmapError::FileUnreadable {
325                path: PathBuf::from("/missing/image.png")
326            }
327        );
328    }
329
330    #[test]
331    fn from_file_status_vendored_returned_null_returns_file_unreadable() {
332        let result = unsafe {
333            super::from_file_status_to_result(
334                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_RETURNED_NULL,
335                std::ptr::null_mut(),
336                std::ptr::null_mut(),
337                "/missing/image.png",
338            )
339        };
340
341        assert_eq!(
342            result.unwrap_err(),
343            MtmdBitmapError::FileUnreadable {
344                path: PathBuf::from("/missing/image.png")
345            }
346        );
347    }
348
349    #[test]
350    fn from_file_status_error_string_allocation_failed_returns_not_enough_memory() {
351        let result = unsafe {
352            super::from_file_status_to_result(
353                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_ERROR_STRING_ALLOCATION_FAILED,
354                std::ptr::null_mut(),
355                std::ptr::null_mut(),
356                "/missing/image.png",
357            )
358        };
359
360        assert_eq!(result.unwrap_err(), MtmdBitmapError::NotEnoughMemory);
361    }
362
363    #[test]
364    fn from_file_status_vendored_threw_cxx_exception_returns_reported() {
365        let result = unsafe {
366            super::from_file_status_to_result(
367                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_THREW_CXX_EXCEPTION,
368                std::ptr::null_mut(),
369                std::ptr::null_mut(),
370                "/missing/image.png",
371            )
372        };
373
374        assert_eq!(
375            result.unwrap_err(),
376            MtmdBitmapError::Reported {
377                message: "unknown error".to_string()
378            }
379        );
380    }
381
382    #[test]
383    #[should_panic(expected = "returned unrecognized status")]
384    fn from_file_status_null_ctx_arg_panics_as_unreachable() {
385        let _result = unsafe {
386            super::from_file_status_to_result(
387                llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_NULL_CTX_ARG,
388                std::ptr::null_mut(),
389                std::ptr::null_mut(),
390                "/missing/image.png",
391            )
392        };
393    }
394}