llama_cpp_bindings/mtmd/
mtmd_bitmap.rs1use 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 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 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 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 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 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}