1use jpegxl_sys::common::memory_manager::{JpegxlAllocFunc, JpegxlFreeFunc, JxlMemoryManager};
21
22#[allow(clippy::module_name_repetitions)]
25pub trait MemoryManager {
26 fn alloc(&self) -> JpegxlAllocFunc;
28 fn free(&self) -> JpegxlFreeFunc;
30
31 #[must_use]
33 fn manager(&self) -> JxlMemoryManager {
34 JxlMemoryManager {
35 opaque: (self as *const Self).cast_mut().cast(),
36 alloc: Some(self.alloc()),
37 free: Some(self.free()),
38 }
39 }
40}
41
42#[cfg(test)]
43pub(crate) mod tests {
44 use std::{
45 ffi::c_void,
46 ptr::null_mut,
47 sync::atomic::{AtomicUsize, Ordering},
48 };
49
50 use testresult::TestResult;
51
52 use crate::{decoder_builder, encoder_builder};
53
54 use super::*;
55 pub struct BumpManager {
57 arena: Vec<u8>,
58 footer: AtomicUsize,
59 }
60
61 impl BumpManager {
62 pub(crate) fn new(n: usize) -> Self {
63 Self {
64 arena: vec![0; n],
65 footer: AtomicUsize::new(0),
66 }
67 }
68 }
69
70 impl MemoryManager for BumpManager {
71 fn alloc(&self) -> JpegxlAllocFunc {
72 #[cfg_attr(coverage_nightly, coverage(off))]
73 unsafe extern "C-unwind" fn alloc(opaque: *mut c_void, size: usize) -> *mut c_void {
74 let mm = &mut *opaque.cast::<BumpManager>();
75
76 let footer = mm.footer.load(Ordering::Acquire);
77 let mut new = footer + size;
78
79 loop {
80 if new > mm.arena.len() {
81 println!("Out of memory");
82 break null_mut();
83 } else if let Err(s) = mm.footer.compare_exchange_weak(
84 footer,
85 new,
86 Ordering::AcqRel,
87 Ordering::Relaxed,
88 ) {
89 new = s + size;
90 } else {
91 let addr = mm.arena.get_unchecked_mut(footer);
92 break (addr as *mut u8).cast();
93 }
94 }
95 }
96
97 alloc
98 }
99
100 fn free(&self) -> JpegxlFreeFunc {
101 #[cfg_attr(coverage_nightly, coverage(off))]
102 unsafe extern "C-unwind" fn free(_opaque: *mut c_void, _address: *mut c_void) {}
103
104 free
105 }
106 }
107 pub struct PanicManager {}
108
109 impl MemoryManager for PanicManager {
110 fn alloc(&self) -> JpegxlAllocFunc {
111 #[cfg_attr(coverage_nightly, coverage(off))]
112 unsafe extern "C-unwind" fn alloc(_opaque: *mut c_void, _size: usize) -> *mut c_void {
113 panic!("Stack unwind test")
114 }
115
116 alloc
117 }
118
119 fn free(&self) -> JpegxlFreeFunc {
120 #[cfg_attr(coverage_nightly, coverage(off))]
121 unsafe extern "C-unwind" fn free(_opaque: *mut c_void, _address: *mut c_void) {
122 debug_assert!(false, "Should not be called");
123 }
124
125 free
126 }
127 }
128
129 #[test]
130 fn test_mm() -> TestResult {
131 let mm = BumpManager::new(1024 * 1024 * 50);
132 let dec = decoder_builder().memory_manager(&mm).build()?;
133 let (meta, img) = dec.decode_with::<u8>(crate::tests::SAMPLE_JXL)?;
134
135 let mut enc = encoder_builder().memory_manager(&mm).build()?;
136 let _ = enc.encode::<u8, u8>(&img, meta.width, meta.height)?;
137
138 Ok(())
139 }
140
141 #[test]
142 #[should_panic = "Stack unwind test"]
143 fn test_unwind() {
144 let mm = PanicManager {};
145 let _ = decoder_builder().memory_manager(&mm).build().unwrap();
146 }
147}