jxl_grid/
lib.rs

1//! This crate provides [`AlignedGrid`], used in various places involving images.
2
3mod alloc_tracker;
4mod mutable_subgrid;
5mod shared_subgrid;
6mod simd;
7pub use alloc_tracker::*;
8pub use mutable_subgrid::*;
9pub use shared_subgrid::*;
10pub use simd::SimdVector;
11
12#[cfg(test)]
13mod test;
14
15/// The error type for failed grid allocation.
16#[derive(Debug)]
17pub struct OutOfMemory {
18    bytes: usize,
19}
20
21impl OutOfMemory {
22    fn new(bytes: usize) -> Self {
23        Self { bytes }
24    }
25
26    /// Returns the size of failed allocation in bytes.
27    pub fn bytes(&self) -> usize {
28        self.bytes
29    }
30}
31
32impl std::error::Error for OutOfMemory {}
33
34impl std::fmt::Display for OutOfMemory {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        write!(f, "failed to allocate {} byte(s)", self.bytes)
37    }
38}
39
40/// A continuous buffer in the "raster order".
41///
42/// The beginning of the buffer is aligned so that it can be used in SIMD instructions.
43pub struct AlignedGrid<S> {
44    width: usize,
45    height: usize,
46    offset: usize,
47    buf: Vec<S>,
48    handle: Option<AllocHandle>,
49}
50
51impl<S> std::fmt::Debug for AlignedGrid<S> {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        f.debug_struct("AlignedGrid")
54            .field("width", &self.width)
55            .field("height", &self.height)
56            .field("offset", &self.offset)
57            .finish_non_exhaustive()
58    }
59}
60
61impl<S> AlignedGrid<S> {
62    /// Creates a zero-sized "empty" grid.
63    #[inline]
64    pub fn empty() -> Self {
65        Self {
66            width: 0,
67            height: 0,
68            offset: 0,
69            buf: Vec::new(),
70            handle: None,
71        }
72    }
73}
74
75impl<S: Default + Clone> AlignedGrid<S> {
76    const ALIGN: usize = 32;
77
78    /// Create a new buffer, recording the allocation if a tracker is given.
79    #[inline]
80    pub fn with_alloc_tracker(
81        width: usize,
82        height: usize,
83        tracker: Option<&AllocTracker>,
84    ) -> Result<Self, OutOfMemory> {
85        let len = width
86            .checked_mul(height)
87            .expect("grid dimensions overflow usize");
88        let buf_len = len
89            .checked_add((Self::ALIGN - 1) / std::mem::size_of::<S>())
90            .expect("aligned grid buffer length overflows usize");
91        let handle = tracker
92            .map(|tracker| tracker.alloc::<S>(buf_len))
93            .transpose()?;
94        let mut buf = vec![S::default(); buf_len];
95
96        let extra = buf.as_ptr() as usize & (Self::ALIGN - 1);
97        let offset = ((Self::ALIGN - extra) % Self::ALIGN) / std::mem::size_of::<S>();
98        let len_with_offset = len
99            .checked_add(offset)
100            .expect("aligned grid buffer length overflows usize");
101        buf.resize_with(len_with_offset, S::default);
102
103        Ok(Self {
104            width,
105            height,
106            offset,
107            buf,
108            handle,
109        })
110    }
111
112    #[inline]
113    fn empty_aligned(
114        width: usize,
115        height: usize,
116        tracker: Option<&AllocTracker>,
117    ) -> Result<Self, OutOfMemory> {
118        let len = width
119            .checked_mul(height)
120            .expect("grid dimensions overflow usize");
121        let buf_len = len
122            .checked_add((Self::ALIGN - 1) / std::mem::size_of::<S>())
123            .expect("aligned grid buffer length overflows usize");
124        let handle = tracker
125            .map(|tracker| tracker.alloc::<S>(buf_len))
126            .transpose()?;
127        let mut buf = Vec::with_capacity(buf_len);
128
129        let extra = buf.as_ptr() as usize & (Self::ALIGN - 1);
130        let offset = ((Self::ALIGN - extra) % Self::ALIGN) / std::mem::size_of::<S>();
131        buf.resize_with(offset, S::default);
132
133        Ok(Self {
134            width,
135            height,
136            offset,
137            buf,
138            handle,
139        })
140    }
141
142    /// Clones the buffer without recording an allocation.
143    pub fn clone_untracked(&self) -> Self {
144        let mut out = Self::empty_aligned(self.width, self.height, None).unwrap();
145        out.buf.extend_from_slice(self.buf());
146        out
147    }
148
149    /// Tries to clone the buffer, and records the allocation in the same tracker as the original
150    /// buffer.
151    pub fn try_clone(&self) -> Result<Self, OutOfMemory> {
152        let mut out = Self::empty_aligned(self.width, self.height, self.tracker().as_ref())?;
153        out.buf.extend_from_slice(self.buf());
154        Ok(out)
155    }
156}
157
158impl<S> AlignedGrid<S> {
159    /// Returns the width of the grid.
160    #[inline]
161    pub fn width(&self) -> usize {
162        self.width
163    }
164
165    /// Returns the height of the grid.
166    #[inline]
167    pub fn height(&self) -> usize {
168        self.height
169    }
170
171    /// Returns allocation tracker associated with the grid.
172    #[inline]
173    pub fn tracker(&self) -> Option<AllocTracker> {
174        self.handle.as_ref().map(|handle| handle.tracker())
175    }
176
177    /// Returns a reference to the sample at the given location.
178    ///
179    /// # Panics
180    /// Panics if the coordinate is out of bounds.
181    #[inline]
182    pub fn get_ref(&self, x: usize, y: usize) -> &S {
183        let width = self.width;
184        let height = self.height;
185        let Some(r) = self.try_get_ref(x, y) else {
186            panic!("coordinate out of range: ({x}, {y}) not in {width}x{height}");
187        };
188
189        r
190    }
191
192    /// Returns a reference to the sample at the given location, or `None` if it is out of bounds.
193    #[inline]
194    pub fn try_get_ref(&self, x: usize, y: usize) -> Option<&S> {
195        if x >= self.width || y >= self.height {
196            return None;
197        }
198
199        Some(&self.buf[y * self.width + x + self.offset])
200    }
201
202    /// Returns a mutable reference to the sample at the given location.
203    ///
204    /// # Panics
205    /// Panics if the coordinate is out of bounds.
206    #[inline]
207    pub fn get_mut(&mut self, x: usize, y: usize) -> &mut S {
208        let width = self.width;
209        let height = self.height;
210        let Some(r) = self.try_get_mut(x, y) else {
211            panic!("coordinate out of range: ({x}, {y}) not in {width}x{height}");
212        };
213
214        r
215    }
216
217    /// Returns a mutable reference to the sample at the given location, or `None` if it is out of
218    /// bounds.
219    #[inline]
220    pub fn try_get_mut(&mut self, x: usize, y: usize) -> Option<&mut S> {
221        if x >= self.width || y >= self.height {
222            return None;
223        }
224
225        Some(&mut self.buf[y * self.width + x + self.offset])
226    }
227
228    /// Returns a slice of a row of samples.
229    ///
230    /// # Panics
231    /// Panics if the row index is out of bounds.
232    #[inline]
233    pub fn get_row(&self, row: usize) -> &[S] {
234        let height = self.height;
235        let Some(slice) = self.try_get_row(row) else {
236            panic!("row index out of range: height is {height} but index is {row}");
237        };
238
239        slice
240    }
241
242    /// Returns a slice of a row of samples, or `None` if it is out of bounds.
243    #[inline]
244    pub fn try_get_row(&self, y: usize) -> Option<&[S]> {
245        if y >= self.height {
246            return None;
247        }
248
249        Some(&self.buf[y * self.width + self.offset..][..self.width])
250    }
251
252    /// Returns a mutable slice of a row of samples.
253    ///
254    /// # Panics
255    /// Panics if the row index is out of bounds.
256    #[inline]
257    pub fn get_row_mut(&mut self, row: usize) -> &mut [S] {
258        let height = self.height;
259        let Some(slice) = self.try_get_row_mut(row) else {
260            panic!("row index out of range: height is {height} but index is {row}");
261        };
262
263        slice
264    }
265
266    /// Returns a mutable slice of a row of samples, or `None` if it is out of bounds.
267    #[inline]
268    pub fn try_get_row_mut(&mut self, y: usize) -> Option<&mut [S]> {
269        if y >= self.height {
270            return None;
271        }
272
273        Some(&mut self.buf[y * self.width + self.offset..][..self.width])
274    }
275
276    /// Returns the immutable slice to the underlying buffer.
277    #[inline]
278    pub fn buf(&self) -> &[S] {
279        &self.buf[self.offset..]
280    }
281
282    /// Returns the mutable slice to the underlying buffer.
283    #[inline]
284    pub fn buf_mut(&mut self) -> &mut [S] {
285        &mut self.buf[self.offset..]
286    }
287
288    /// Borrows the grid into a `SharedSubgrid`.
289    #[inline]
290    pub fn as_subgrid(&self) -> SharedSubgrid<'_, S> {
291        SharedSubgrid::from(self)
292    }
293
294    /// Borrows the grid into a `MutableSubgrid`.
295    #[inline]
296    pub fn as_subgrid_mut(&mut self) -> MutableSubgrid<'_, S> {
297        MutableSubgrid::from(self)
298    }
299}
300
301impl<V: Copy> AlignedGrid<V> {
302    /// Returns a copy of sample at the given location.
303    ///
304    /// # Panics
305    /// Panics if the coordinate is out of range.
306    #[inline]
307    pub fn get(&self, x: usize, y: usize) -> V {
308        *self.get_ref(x, y)
309    }
310}
jxl_grid/lib.rs

jxl_grid/
lib.rs
