Skip to main content

llama_cpp_bindings/mtmd/
mtmd_bitmap.rs

1use std::ffi::{CStr, CString, c_char};
2use std::ptr::NonNull;
3use std::slice;
4
5use super::mtmd_context::MtmdContext;
6use super::mtmd_error::MtmdBitmapError;
7
8fn cstr_ptr_to_optional_string(ptr: *const c_char) -> Option<String> {
9    if ptr.is_null() {
10        None
11    } else {
12        let id = unsafe { CStr::from_ptr(ptr) }
13            .to_string_lossy()
14            .into_owned();
15
16        Some(id)
17    }
18}
19
20/// Safe wrapper around `mtmd_bitmap`.
21///
22/// Represents bitmap data for images or audio that can be processed
23/// by the multimodal system. For images, data is stored in RGB format.
24/// For audio, data is stored as PCM F32 samples.
25#[derive(Debug, Clone)]
26pub struct MtmdBitmap {
27    /// Raw pointer to the underlying `mtmd_bitmap`.
28    pub bitmap: NonNull<llama_cpp_bindings_sys::mtmd_bitmap>,
29}
30
31unsafe impl Send for MtmdBitmap {}
32unsafe impl Sync for MtmdBitmap {}
33
34impl MtmdBitmap {
35    /// Create a bitmap from image data in RGB format.
36    ///
37    /// # Errors
38    ///
39    /// * `InvalidDataSize` - Data length doesn't match `nx * ny * 3`
40    /// * `NullResult` - Underlying C function returned null
41    ///
42    /// # Examples
43    ///
44    /// ```
45    /// use llama_cpp_bindings::mtmd::MtmdBitmap;
46    ///
47    /// // Create a 2x2 red image
48    /// let red_pixel = [255, 0, 0]; // RGB values for red
49    /// let image_data = red_pixel.repeat(4); // 2x2 = 4 pixels
50    ///
51    /// let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data);
52    /// assert!(bitmap.is_ok());
53    /// ```
54    pub fn from_image_data(nx: u32, ny: u32, data: &[u8]) -> Result<Self, MtmdBitmapError> {
55        if nx < 2 || ny < 2 {
56            return Err(MtmdBitmapError::ImageDimensionsTooSmall(nx, ny));
57        }
58
59        if data.len() != (nx * ny * 3) as usize {
60            return Err(MtmdBitmapError::InvalidDataSize);
61        }
62
63        let bitmap = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_init(nx, ny, data.as_ptr()) };
64
65        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::NullResult)?;
66
67        Ok(Self { bitmap })
68    }
69
70    /// Create a bitmap from audio data in PCM F32 format.
71    ///
72    /// # Errors
73    ///
74    /// * `NullResult` - Underlying C function returned null
75    ///
76    /// # Examples
77    ///
78    /// ```
79    /// use llama_cpp_bindings::mtmd::MtmdBitmap;
80    ///
81    /// // Create a simple sine wave audio sample
82    /// let audio_data: Vec<f32> = (0..100)
83    ///     .map(|sample_index| (sample_index as f32 * 0.1).sin())
84    ///     .collect();
85    ///
86    /// let bitmap = MtmdBitmap::from_audio_data(&audio_data);
87    /// // Note: This will likely fail without proper MTMD context setup
88    /// ```
89    pub fn from_audio_data(data: &[f32]) -> Result<Self, MtmdBitmapError> {
90        let bitmap = unsafe {
91            llama_cpp_bindings_sys::mtmd_bitmap_init_from_audio(data.len(), data.as_ptr())
92        };
93
94        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::NullResult)?;
95
96        Ok(Self { bitmap })
97    }
98
99    /// Create a bitmap from a file.
100    ///
101    /// Supported formats:
102    /// - Images: formats supported by `stb_image` (jpg, png, bmp, gif, etc.)
103    /// - Audio: formats supported by miniaudio (wav, mp3, flac)
104    ///
105    /// # Errors
106    ///
107    /// * `CStringError` - Path contains null bytes
108    /// * `NullResult` - File could not be loaded or processed
109    pub fn from_file(ctx: &MtmdContext, path: &str) -> Result<Self, MtmdBitmapError> {
110        let path_cstr = CString::new(path)?;
111        let bitmap = unsafe {
112            llama_cpp_bindings_sys::mtmd_helper_bitmap_init_from_file(
113                ctx.context.as_ptr(),
114                path_cstr.as_ptr(),
115            )
116        };
117
118        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::NullResult)?;
119
120        Ok(Self { bitmap })
121    }
122
123    /// Create a bitmap from a buffer containing file data.
124    ///
125    /// Supported formats:
126    /// - Images: formats supported by `stb_image` (jpg, png, bmp, gif, etc.)
127    /// - Audio: formats supported by miniaudio (wav, mp3, flac)
128    ///
129    /// # Errors
130    ///
131    /// * `NullResult` - Buffer could not be processed
132    pub fn from_buffer(ctx: &MtmdContext, data: &[u8]) -> Result<Self, MtmdBitmapError> {
133        let bitmap = unsafe {
134            llama_cpp_bindings_sys::mtmd_helper_bitmap_init_from_buf(
135                ctx.context.as_ptr(),
136                data.as_ptr(),
137                data.len(),
138            )
139        };
140
141        let bitmap = NonNull::new(bitmap).ok_or(MtmdBitmapError::NullResult)?;
142
143        Ok(Self { bitmap })
144    }
145
146    /// Get bitmap width in pixels.
147    #[must_use]
148    pub fn nx(&self) -> u32 {
149        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_nx(self.bitmap.as_ptr()) }
150    }
151
152    /// Get bitmap height in pixels.
153    #[must_use]
154    pub fn ny(&self) -> u32 {
155        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_ny(self.bitmap.as_ptr()) }
156    }
157
158    /// Get bitmap data as a byte slice.
159    ///
160    /// For images: RGB format with length `nx * ny * 3`
161    /// For audio: PCM F32 format with length `n_samples * 4`
162    #[must_use]
163    pub fn data(&self) -> &[u8] {
164        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_data(self.bitmap.as_ptr()) };
165        let len = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_n_bytes(self.bitmap.as_ptr()) };
166        unsafe { slice::from_raw_parts(ptr, len) }
167    }
168
169    /// Check if this bitmap contains audio data (vs image data).
170    #[must_use]
171    pub fn is_audio(&self) -> bool {
172        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_is_audio(self.bitmap.as_ptr()) }
173    }
174
175    /// Get the bitmap's optional ID string.
176    #[must_use]
177    pub fn id(&self) -> Option<String> {
178        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_bitmap_get_id(self.bitmap.as_ptr()) };
179
180        cstr_ptr_to_optional_string(ptr)
181    }
182
183    /// Set the bitmap's ID string.
184    ///
185    /// # Errors
186    ///
187    /// Returns an error if the ID string contains null bytes.
188    ///
189    /// # Examples
190    ///
191    /// ```no_run
192    /// # use llama_cpp_bindings::mtmd::MtmdBitmap;
193    /// # fn example(bitmap: &MtmdBitmap) -> Result<(), Box<dyn std::error::Error>> {
194    /// bitmap.set_id("image_001")?;
195    /// assert_eq!(bitmap.id(), Some("image_001".to_string()));
196    /// # Ok(())
197    /// # }
198    /// ```
199    pub fn set_id(&self, id: &str) -> Result<(), std::ffi::NulError> {
200        let id_cstr = CString::new(id)?;
201        unsafe {
202            llama_cpp_bindings_sys::mtmd_bitmap_set_id(self.bitmap.as_ptr(), id_cstr.as_ptr());
203        }
204
205        Ok(())
206    }
207}
208
209impl Drop for MtmdBitmap {
210    fn drop(&mut self) {
211        unsafe { llama_cpp_bindings_sys::mtmd_bitmap_free(self.bitmap.as_ptr()) }
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::MtmdBitmap;
218    use super::MtmdBitmapError;
219
220    #[test]
221    fn cstr_ptr_to_optional_string_returns_none_for_null() {
222        assert!(super::cstr_ptr_to_optional_string(std::ptr::null()).is_none());
223    }
224
225    #[test]
226    fn cstr_ptr_to_optional_string_returns_some_for_valid() {
227        let cstr = std::ffi::CString::new("hello").unwrap();
228        let result = super::cstr_ptr_to_optional_string(cstr.as_ptr());
229
230        assert_eq!(result, Some("hello".to_string()));
231    }
232
233    #[test]
234    fn from_image_data_creates_valid_bitmap() {
235        let red_pixel: [u8; 3] = [255, 0, 0];
236        let image_data: Vec<u8> = red_pixel.repeat(4);
237        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
238        assert_eq!(bitmap.nx(), 2);
239        assert_eq!(bitmap.ny(), 2);
240        assert!(!bitmap.is_audio());
241    }
242
243    #[test]
244    fn invalid_data_size_returns_error() {
245        let too_short = vec![0u8; 5];
246        let result = MtmdBitmap::from_image_data(2, 2, &too_short);
247        assert!(result.is_err());
248    }
249
250    #[test]
251    fn from_image_data_rejects_dimensions_below_minimum() {
252        let result_1x1 = MtmdBitmap::from_image_data(1, 1, &[0u8; 3]);
253        let result_1x2 = MtmdBitmap::from_image_data(1, 2, &[0u8; 6]);
254        let result_2x1 = MtmdBitmap::from_image_data(2, 1, &[0u8; 6]);
255        let result_0x0 = MtmdBitmap::from_image_data(0, 0, &[]);
256
257        assert!(matches!(
258            result_1x1,
259            Err(MtmdBitmapError::ImageDimensionsTooSmall(1, 1))
260        ));
261        assert!(matches!(
262            result_1x2,
263            Err(MtmdBitmapError::ImageDimensionsTooSmall(1, 2))
264        ));
265        assert!(matches!(
266            result_2x1,
267            Err(MtmdBitmapError::ImageDimensionsTooSmall(2, 1))
268        ));
269        assert!(matches!(
270            result_0x0,
271            Err(MtmdBitmapError::ImageDimensionsTooSmall(0, 0))
272        ));
273    }
274
275    #[test]
276    fn set_id_changes_id() {
277        let image_data = vec![0u8; 12];
278        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
279        bitmap.set_id("test_image").unwrap();
280
281        assert_eq!(bitmap.id().as_deref(), Some("test_image"));
282    }
283
284    #[test]
285    fn from_audio_data_creates_valid_bitmap() {
286        #[expect(
287            clippy::cast_precision_loss,
288            reason = "test fixture casts a small i32 (0..100) to f32 to synthesise a sine wave; \
289                      the values are well within f32's exact-representation range"
290        )]
291        let audio_samples: Vec<f32> = (0..100).map(|index| (index as f32 * 0.1).sin()).collect();
292        let bitmap = MtmdBitmap::from_audio_data(&audio_samples).unwrap();
293
294        assert!(bitmap.is_audio());
295    }
296
297    #[test]
298    fn data_returns_expected_bytes_for_image() {
299        let pixel_data: Vec<u8> = vec![255, 0, 0, 0, 255, 0, 0, 0, 255, 128, 128, 128];
300        let bitmap = MtmdBitmap::from_image_data(2, 2, &pixel_data).unwrap();
301        let returned_data = bitmap.data();
302
303        assert_eq!(returned_data, &pixel_data);
304    }
305
306    #[test]
307    fn id_returns_some_by_default() {
308        let image_data = vec![0u8; 12];
309        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
310
311        assert!(bitmap.id().is_some());
312    }
313
314    #[test]
315    fn id_returns_custom_value_after_set() {
316        let image_data = vec![0u8; 12];
317        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
318        bitmap.set_id("my_image").unwrap();
319
320        assert_eq!(bitmap.id(), Some("my_image".to_string()));
321    }
322
323    #[test]
324    fn set_id_with_null_byte_returns_error() {
325        let image_data = vec![0u8; 12];
326        let bitmap = MtmdBitmap::from_image_data(2, 2, &image_data).unwrap();
327        let result = bitmap.set_id("id\0null");
328
329        assert!(result.is_err());
330    }
331}