jpegxl_rs/
memory.rs

1/*
2 * This file is part of jpegxl-rs.
3 *
4 * jpegxl-rs is free software: you can redistribute it and/or modify
5 * it under the terms of the GNU General Public License as published by
6 * the Free Software Foundation, either version 3 of the License, or
7 * (at your option) any later version.
8 *
9 * jpegxl-rs is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12 * GNU General Public License for more details.
13 *
14 * You should have received a copy of the GNU General Public License
15 * along with jpegxl-rs.  If not, see <https://www.gnu.org/licenses/>.
16 */
17
18//! Memory manager interface
19
20use jpegxl_sys::common::memory_manager::{JpegxlAllocFunc, JpegxlFreeFunc, JxlMemoryManager};
21
22/// General trait for a memory manager
23
24#[allow(clippy::module_name_repetitions)]
25pub trait MemoryManager {
26    /// Return a custom allocating function
27    fn alloc(&self) -> JpegxlAllocFunc;
28    /// Return a custom deallocating function
29    fn free(&self) -> JpegxlFreeFunc;
30
31    /// Helper conversion function for C API
32    #[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    /// Example implementation of [`MemoryManager`] of a fixed size allocator
56    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}