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#[derive(Debug, Clone)]
24pub struct MtmdBitmap {
25    pub bitmap: NonNull<llama_cpp_bindings_sys::mtmd_bitmap>,
26}
27
28unsafe impl Send for MtmdBitmap {}
29unsafe impl Sync for MtmdBitmap {}
30
31impl MtmdBitmap {
32    /// # Errors
33    ///
34    /// * `InvalidDataSize` - Data length doesn't match `nx * ny * 3`
35    /// * `NullResult` - Underlying C function returned null
36    ///
37    pub fn from_image_data(nx: u32, ny: u32, data: &[u8]) -> Result<Self, MtmdBitmapError> {
38        if nx < 2 || ny < 2 {
39            return Err(MtmdBitmapError::ImageDimensionsTooSmall(nx, ny));
40        }
41
42        if data.len() != (nx * ny * 3) as usize {
43            return Err(MtmdBitmapError::InvalidDataSize);
44        }
45
46        let bitmap = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_init(nx, ny, data.as_ptr()) };
47
48        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::BitmapDecodeFailed)?;
49
50        Ok(Self { bitmap })
51    }
52
53    /// # Errors
54    ///
55    /// * `NullResult` - Underlying C function returned null
56    ///
57    pub fn from_audio_data(data: &[f32]) -> Result<Self, MtmdBitmapError> {
58        let bitmap = unsafe {
59            llama_cpp_bindings_sys::mtmd_bitmap_init_from_audio(data.len(), data.as_ptr())
60        };
61
62        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::BitmapDecodeFailed)?;
63
64        Ok(Self { bitmap })
65    }
66
67    /// # Errors
68    ///
69    /// Returns an [`MtmdBitmapError`] variant matching the wrapper's status code.
70    pub fn from_file(ctx: &MtmdContext, path: &str) -> Result<Self, MtmdBitmapError> {
71        let path_cstr = CString::new(path)?;
72        let mut out_bitmap: *mut llama_cpp_bindings_sys::mtmd_bitmap = std::ptr::null_mut();
73        let mut out_error: *mut c_char = std::ptr::null_mut();
74
75        let status = unsafe {
76            llama_cpp_bindings_sys::llama_rs_mtmd_bitmap_init_from_file(
77                ctx.context.as_ptr(),
78                path_cstr.as_ptr(),
79                &raw mut out_bitmap,
80                &raw mut out_error,
81            )
82        };
83
84        match status {
85            llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_OK => {
86                let bitmap = NonNull::new(out_bitmap).ok_or_else(|| {
87                    MtmdBitmapError::FileUnreadable {
88                        path: PathBuf::from(path),
89                    }
90                })?;
91                Ok(Self { bitmap })
92            }
93            llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_RETURNED_NULL => {
94                Err(MtmdBitmapError::FileUnreadable {
95                    path: PathBuf::from(path),
96                })
97            }
98            llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_ERROR_STRING_ALLOCATION_FAILED => {
99                Err(MtmdBitmapError::NotEnoughMemory)
100            }
101            llama_cpp_bindings_sys::LLAMA_RS_MTMD_BITMAP_INIT_FROM_FILE_VENDORED_THREW_CXX_EXCEPTION => {
102                let message = unsafe { read_and_free_cpp_error(out_error) };
103                Err(MtmdBitmapError::Reported { message })
104            }
105            other => unreachable!(
106                "llama_rs_mtmd_bitmap_init_from_file returned unrecognized status: {other}"
107            ),
108        }
109    }
110
111    /// # Errors
112    ///
113    /// * `NullResult` - Buffer could not be processed
114    pub fn from_buffer(ctx: &MtmdContext, data: &[u8]) -> Result<Self, MtmdBitmapError> {
115        let bitmap = unsafe {
116            llama_cpp_bindings_sys::mtmd_helper_bitmap_init_from_buf(
117                ctx.context.as_ptr(),
118                data.as_ptr(),
119                data.len(),
120            )
121        };
122
123        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::BitmapDecodeFailed)?;
124
125        Ok(Self { bitmap })
126    }
127
128    #[must_use]
129    pub fn nx(&self) -> u32 {
130        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_nx(self.bitmap.as_ptr()) }
131    }
132
133    #[must_use]
134    pub fn ny(&self) -> u32 {
135        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_ny(self.bitmap.as_ptr()) }
136    }
137
138    #[must_use]
139    pub fn data(&self) -> &[u8] {
140        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_data(self.bitmap.as_ptr()) };
141        let len = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_n_bytes(self.bitmap.as_ptr()) };
142        unsafe { slice::from_raw_parts(ptr, len) }
143    }
144
145    #[must_use]
146    pub fn is_audio(&self) -> bool {
147        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_is_audio(self.bitmap.as_ptr()) }
148    }
149
150    #[must_use]
151    pub fn id(&self) -> Option<String> {
152        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_id(self.bitmap.as_ptr()) };
153
154        cstr_ptr_to_optional_string(ptr)
155    }
156
157    /// # Errors
158    ///
159    /// Returns an error if the ID string contains null bytes.
160    ///
161    pub fn set_id(&self, id: &str) -> Result<(), std::ffi::NulError> {
162        let id_cstr = CString::new(id)?;
163        unsafe {
164            llama_cpp_bindings_sys::mtmd_bitmap_set_id(self.bitmap.as_ptr(), id_cstr.as_ptr());
165        }
166
167        Ok(())
168    }
169}
170
171impl Drop for MtmdBitmap {
172    fn drop(&mut self) {
173        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_free(self.bitmap.as_ptr()) }
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::MtmdBitmap;
180    use super::MtmdBitmapError;
181
182    #[test]
183    fn cstr_ptr_to_optional_string_returns_none_for_null() {
184        assert!(super::cstr_ptr_to_optional_string(std::ptr::null()).is_none());
185    }
186
187    #[test]
188    fn cstr_ptr_to_optional_string_returns_some_for_valid() {
189        let cstr = std::ffi::CString::new("hello").unwrap();
190        let result = super::cstr_ptr_to_optional_string(cstr.as_ptr());
191
192        assert_eq!(result, Some("hello".to_string()));
193    }
194
195    #[test]
196    fn from_image_data_creates_valid_bitmap() {
197        let red_pixel: [u8; 3] = [255, 0, 0];
198        let image_data: Vec<u8> = red_pixel.repeat(4);
199        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
200        assert_eq!(bitmap.nx(), 2);
201        assert_eq!(bitmap.ny(), 2);
202        assert!(!bitmap.is_audio());
203    }
204
205    #[test]
206    fn invalid_data_size_returns_error() {
207        let too_short = vec![0u8; 5];
208        let result = MtmdBitmap::from_image_data(2, 2, &too_short);
209        assert!(result.is_err());
210    }
211
212    #[test]
213    fn from_image_data_rejects_dimensions_below_minimum() {
214        let result_1x1 = MtmdBitmap::from_image_data(1, 1, &[0u8; 3]);
215        let result_1x2 = MtmdBitmap::from_image_data(1, 2, &[0u8; 6]);
216        let result_2x1 = MtmdBitmap::from_image_data(2, 1, &[0u8; 6]);
217        let result_0x0 = MtmdBitmap::from_image_data(0, 0, &[]);
218
219        assert!(matches!(
220            result_1x1,
221            Err(MtmdBitmapError::ImageDimensionsTooSmall(1, 1))
222        ));
223        assert!(matches!(
224            result_1x2,
225            Err(MtmdBitmapError::ImageDimensionsTooSmall(1, 2))
226        ));
227        assert!(matches!(
228            result_2x1,
229            Err(MtmdBitmapError::ImageDimensionsTooSmall(2, 1))
230        ));
231        assert!(matches!(
232            result_0x0,
233            Err(MtmdBitmapError::ImageDimensionsTooSmall(0, 0))
234        ));
235    }
236
237    #[test]
238    fn set_id_changes_id() {
239        let image_data = vec![0u8; 12];
240        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
241        bitmap.set_id("test_image").unwrap();
242
243        assert_eq!(bitmap.id().as_deref(), Some("test_image"));
244    }
245
246    #[test]
247    fn from_audio_data_creates_valid_bitmap() {
248        #[expect(
249            clippy::cast_precision_loss,
250            reason = "test fixture casts a small i32 (0..100) to f32 to synthesise a sine wave; \
251                      the values are well within f32's exact-representation range"
252        )]
253        let audio_samples: Vec<f32> = (0..100).map(|index| (index as f32 * 0.1).sin()).collect();
254        let bitmap = MtmdBitmap::from_audio_data(&audio_samples).unwrap();
255
256        assert!(bitmap.is_audio());
257    }
258
259    #[test]
260    fn data_returns_expected_bytes_for_image() {
261        let pixel_data: Vec<u8> = vec![255, 0, 0, 0, 255, 0, 0, 0, 255, 128, 128, 128];
262        let bitmap = MtmdBitmap::from_image_data(2, 2, &pixel_data).unwrap();
263        let returned_data = bitmap.data();
264
265        assert_eq!(returned_data, &pixel_data);
266    }
267
268    #[test]
269    fn id_returns_some_by_default() {
270        let image_data = vec![0u8; 12];
271        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
272
273        assert!(bitmap.id().is_some());
274    }
275
276    #[test]
277    fn id_returns_custom_value_after_set() {
278        let image_data = vec![0u8; 12];
279        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
280        bitmap.set_id("my_image").unwrap();
281
282        assert_eq!(bitmap.id(), Some("my_image".to_string()));
283    }
284
285    #[test]
286    fn set_id_with_null_byte_returns_error() {
287        let image_data = vec![0u8; 12];
288        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
289        let result = bitmap.set_id("id\0null");
290
291        assert!(result.is_err());
292    }
293}