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
23unsafe 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 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 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 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 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 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}