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