1#[cfg(any(feature = "std", feature = "alloc"))]
2use alloc::{vec, vec::Vec};
3#[cfg(feature = "std")]
4use std::io::Read;
5
6// TODO: can be removed once https://github.com/rust-lang/rust/issues/74985 is stable
7use bytemuck::{cast_slice_mut, Pod};
8
9use crate::consts::{
10 QOI_HEADER_SIZE, QOI_OP_DIFF, QOI_OP_INDEX, QOI_OP_LUMA, QOI_OP_RGB, QOI_OP_RGBA, QOI_OP_RUN,
11 QOI_PADDING, QOI_PADDING_SIZE,
12};
13use crate::error::{Error, Result};
14use crate::header::Header;
15use crate::pixel::{Pixel, SupportedChannels};
16use crate::types::Channels;
17use crate::utils::{cold, unlikely};
18
19const QOI_OP_INDEX_END: u8 = QOI_OP_INDEX | 0x3f;
20const QOI_OP_RUN_END: u8 = QOI_OP_RUN | 0x3d; // <- note, 0x3d (not 0x3f)
21const QOI_OP_DIFF_END: u8 = QOI_OP_DIFF | 0x3f;
22const QOI_OP_LUMA_END: u8 = QOI_OP_LUMA | 0x3f;
23
24#[inline]
25fn decode_impl_slice<const N: usize, const RGBA: bool>(data: &[u8], out: &mut [u8]) -> Result<usize>
26where
27 Pixel<N>: SupportedChannels,
28 [u8; N]: Pod,
29{
30 let mut pixels = cast_slice_mut::<_, [u8; N]>(out);
31 let data_len = data.len();
32 let mut data = data;
33
34 let mut index = [Pixel::<4>::new(); 256];
35 let mut px = Pixel::<N>::new().with_a(0xff);
36 let mut px_rgba: Pixel<4>;
37
38 while let [px_out, ptail @ ..] = pixels {
39 pixels = ptail;
40 match data {
41 [b1 @ QOI_OP_INDEX..=QOI_OP_INDEX_END, dtail @ ..] => {
42 px_rgba = index[*b1 as usize];
43 px.update(px_rgba);
44 *px_out = px.into();
45 data = dtail;
46 continue;
47 }
48 [QOI_OP_RGB, r, g, b, dtail @ ..] => {
49 px.update_rgb(*r, *g, *b);
50 data = dtail;
51 }
52 [QOI_OP_RGBA, r, g, b, a, dtail @ ..] if RGBA => {
53 px.update_rgba(*r, *g, *b, *a);
54 data = dtail;
55 }
56 [b1 @ QOI_OP_RUN..=QOI_OP_RUN_END, dtail @ ..] => {
57 *px_out = px.into();
58 let run = ((b1 & 0x3f) as usize).min(pixels.len());
59 let (phead, ptail) = pixels.split_at_mut(run); // can't panic
60 phead.fill(px.into());
61 pixels = ptail;
62 data = dtail;
63 continue;
64 }
65 [b1 @ QOI_OP_DIFF..=QOI_OP_DIFF_END, dtail @ ..] => {
66 px.update_diff(*b1);
67 data = dtail;
68 }
69 [b1 @ QOI_OP_LUMA..=QOI_OP_LUMA_END, b2, dtail @ ..] => {
70 px.update_luma(*b1, *b2);
71 data = dtail;
72 }
73 _ => {
74 cold();
75 if unlikely(data.len() < QOI_PADDING_SIZE) {
76 return Err(Error::UnexpectedBufferEnd);
77 }
78 }
79 }
80
81 px_rgba = px.as_rgba(0xff);
82 index[px_rgba.hash_index() as usize] = px_rgba;
83 *px_out = px.into();
84 }
85
86 if unlikely(data.len() < QOI_PADDING_SIZE) {
87 return Err(Error::UnexpectedBufferEnd);
88 } else if unlikely(data[..QOI_PADDING_SIZE] != QOI_PADDING) {
89 return Err(Error::InvalidPadding);
90 }
91
92 Ok(data_len.saturating_sub(data.len()).saturating_sub(QOI_PADDING_SIZE))
93}
94
95#[inline]
96fn decode_impl_slice_all(
97 data: &[u8], out: &mut [u8], channels: u8, src_channels: u8,
98) -> Result<usize> {
99 match (channels, src_channels) {
100 (3, 3) => decode_impl_slice::<3, false>(data, out),
101 (3, 4) => decode_impl_slice::<3, true>(data, out),
102 (4, 3) => decode_impl_slice::<4, false>(data, out),
103 (4, 4) => decode_impl_slice::<4, true>(data, out),
104 _ => {
105 cold();
106 Err(Error::InvalidChannels { channels })
107 }
108 }
109}
110
111/// Decode the image into a pre-allocated buffer.
112///
113/// Note: the resulting number of channels will match the header. In order to change
114/// the number of channels, use [`Decoder::with_channels`].
115#[inline]
116pub fn decode_to_buf(buf: impl AsMut<[u8]>, data: impl AsRef<[u8]>) -> Result<Header> {
117 let mut decoder: Decoder> = Decoder::new(&data)?;
118 decoder.decode_to_buf(buf)?;
119 Ok(*decoder.header())
120}
121
122/// Decode the image into a newly allocated vector.
123///
124/// Note: the resulting number of channels will match the header. In order to change
125/// the number of channels, use [`Decoder::with_channels`].
126#[cfg(any(feature = "std", feature = "alloc"))]
127#[inline]
128pub fn decode_to_vec(data: impl AsRef<[u8]>) -> Result<(Header, Vec<u8>)> {
129 let mut decoder: Decoder> = Decoder::new(&data)?;
130 let out: Vec = decoder.decode_to_vec()?;
131 Ok((*decoder.header(), out))
132}
133
134/// Decode the image header from a slice of bytes.
135#[inline]
136pub fn decode_header(data: impl AsRef<[u8]>) -> Result<Header> {
137 Header::decode(data)
138}
139
140#[cfg(any(feature = "std"))]
141#[inline]
142fn decode_impl_stream<R: Read, const N: usize, const RGBA: bool>(
143 data: &mut R, out: &mut [u8],
144) -> Result<()>
145where
146 Pixel<N>: SupportedChannels,
147 [u8; N]: Pod,
148{
149 let mut pixels = cast_slice_mut::<_, [u8; N]>(out);
150
151 let mut index = [Pixel::<N>::new(); 256];
152 let mut px = Pixel::<N>::new().with_a(0xff);
153
154 while let [px_out, ptail @ ..] = pixels {
155 pixels = ptail;
156 let mut p = [0];
157 data.read_exact(&mut p)?;
158 let [b1] = p;
159 match b1 {
160 QOI_OP_INDEX..=QOI_OP_INDEX_END => {
161 px = index[b1 as usize];
162 *px_out = px.into();
163 continue;
164 }
165 QOI_OP_RGB => {
166 let mut p = [0; 3];
167 data.read_exact(&mut p)?;
168 px.update_rgb(p[0], p[1], p[2]);
169 }
170 QOI_OP_RGBA if RGBA => {
171 let mut p = [0; 4];
172 data.read_exact(&mut p)?;
173 px.update_rgba(p[0], p[1], p[2], p[3]);
174 }
175 QOI_OP_RUN..=QOI_OP_RUN_END => {
176 *px_out = px.into();
177 let run = ((b1 & 0x3f) as usize).min(pixels.len());
178 let (phead, ptail) = pixels.split_at_mut(run); // can't panic
179 phead.fill(px.into());
180 pixels = ptail;
181 continue;
182 }
183 QOI_OP_DIFF..=QOI_OP_DIFF_END => {
184 px.update_diff(b1);
185 }
186 QOI_OP_LUMA..=QOI_OP_LUMA_END => {
187 let mut p = [0];
188 data.read_exact(&mut p)?;
189 let [b2] = p;
190 px.update_luma(b1, b2);
191 }
192 _ => {
193 cold();
194 }
195 }
196
197 index[px.hash_index() as usize] = px;
198 *px_out = px.into();
199 }
200
201 let mut p = [0_u8; QOI_PADDING_SIZE];
202 data.read_exact(&mut p)?;
203 if unlikely(p != QOI_PADDING) {
204 return Err(Error::InvalidPadding);
205 }
206
207 Ok(())
208}
209
210#[cfg(feature = "std")]
211#[inline]
212fn decode_impl_stream_all<R: Read>(
213 data: &mut R, out: &mut [u8], channels: u8, src_channels: u8,
214) -> Result<()> {
215 match (channels, src_channels) {
216 (3, 3) => decode_impl_stream::<_, 3, false>(data, out),
217 (3, 4) => decode_impl_stream::<_, 3, true>(data, out),
218 (4, 3) => decode_impl_stream::<_, 4, false>(data, out),
219 (4, 4) => decode_impl_stream::<_, 4, true>(data, out),
220 _ => {
221 cold();
222 Err(Error::InvalidChannels { channels })
223 }
224 }
225}
226
227#[doc(hidden)]
228pub trait Reader: Sized {
229 fn decode_header(&mut self) -> Result<Header>;
230 fn decode_image(&mut self, out: &mut [u8], channels: u8, src_channels: u8) -> Result<()>;
231}
232
233pub struct Bytes<'a>(&'a [u8]);
234
235impl<'a> Bytes<'a> {
236 #[inline]
237 pub const fn new(buf: &'a [u8]) -> Self {
238 Self(buf)
239 }
240
241 #[inline]
242 pub const fn as_slice(&self) -> &[u8] {
243 self.0
244 }
245}
246
247impl<'a> Reader for Bytes<'a> {
248 #[inline]
249 fn decode_header(&mut self) -> Result<Header> {
250 let header: Header = Header::decode(self.0)?;
251 self.0 = &self.0[QOI_HEADER_SIZE..]; // can't panic
252 Ok(header)
253 }
254
255 #[inline]
256 fn decode_image(&mut self, out: &mut [u8], channels: u8, src_channels: u8) -> Result<()> {
257 let n_read: usize = decode_impl_slice_all(self.0, out, channels, src_channels)?;
258 self.0 = &self.0[n_read..];
259 Ok(())
260 }
261}
262
263#[cfg(feature = "std")]
264impl<R: Read> Reader for R {
265 #[inline]
266 fn decode_header(&mut self) -> Result<Header> {
267 let mut b: [u8; 14] = [0; QOI_HEADER_SIZE];
268 self.read_exact(&mut b)?;
269 Header::decode(data:b)
270 }
271
272 #[inline]
273 fn decode_image(&mut self, out: &mut [u8], channels: u8, src_channels: u8) -> Result<()> {
274 decode_impl_stream_all(self, out, channels, src_channels)
275 }
276}
277
278/// Decode QOI images from slices or from streams.
279#[derive(Clone)]
280pub struct Decoder<R> {
281 reader: R,
282 header: Header,
283 channels: Channels,
284}
285
286impl<'a> Decoder<Bytes<'a>> {
287 /// Creates a new decoder from a slice of bytes.
288 ///
289 /// The header will be decoded immediately upon construction.
290 ///
291 /// Note: this provides the most efficient decoding, but requires the source data to
292 /// be loaded in memory in order to decode it. In order to decode from a generic
293 /// stream, use [`Decoder::from_stream`] instead.
294 #[inline]
295 pub fn new(data: &'a (impl AsRef<[u8]> + ?Sized)) -> Result<Self> {
296 Self::new_impl(reader:Bytes::new(buf:data.as_ref()))
297 }
298
299 /// Returns the undecoded tail of the input slice of bytes.
300 #[inline]
301 pub const fn data(&self) -> &[u8] {
302 self.reader.as_slice()
303 }
304}
305
306#[cfg(feature = "std")]
307impl<R: Read> Decoder<R> {
308 /// Creates a new decoder from a generic reader that implements [`Read`](std::io::Read).
309 ///
310 /// The header will be decoded immediately upon construction.
311 ///
312 /// Note: while it's possible to pass a `&[u8]` slice here since it implements `Read`, it
313 /// would be more efficient to use a specialized constructor instead: [`Decoder::new`].
314 #[inline]
315 pub fn from_stream(reader: R) -> Result<Self> {
316 Self::new_impl(reader)
317 }
318
319 /// Returns an immutable reference to the underlying reader.
320 #[inline]
321 pub const fn reader(&self) -> &R {
322 &self.reader
323 }
324
325 /// Consumes the decoder and returns the underlying reader back.
326 #[inline]
327 #[allow(clippy::missing_const_for_fn)]
328 pub fn into_reader(self) -> R {
329 self.reader
330 }
331}
332
333impl<R: Reader> Decoder<R> {
334 #[inline]
335 fn new_impl(mut reader: R) -> Result<Self> {
336 let header = reader.decode_header()?;
337 Ok(Self { reader, header, channels: header.channels })
338 }
339
340 /// Returns a new decoder with modified number of channels.
341 ///
342 /// By default, the number of channels in the decoded image will be equal
343 /// to whatever is specified in the header. However, it is also possible
344 /// to decode RGB into RGBA (in which case the alpha channel will be set
345 /// to 255), and vice versa (in which case the alpha channel will be ignored).
346 #[inline]
347 pub const fn with_channels(mut self, channels: Channels) -> Self {
348 self.channels = channels;
349 self
350 }
351
352 /// Returns the number of channels in the decoded image.
353 ///
354 /// Note: this may differ from the number of channels specified in the header.
355 #[inline]
356 pub const fn channels(&self) -> Channels {
357 self.channels
358 }
359
360 /// Returns the decoded image header.
361 #[inline]
362 pub const fn header(&self) -> &Header {
363 &self.header
364 }
365
366 /// The number of bytes the decoded image will take.
367 ///
368 /// Can be used to pre-allocate the buffer to decode the image into.
369 #[inline]
370 pub const fn required_buf_len(&self) -> usize {
371 self.header.n_pixels().saturating_mul(self.channels.as_u8() as usize)
372 }
373
374 /// Decodes the image to a pre-allocated buffer and returns the number of bytes written.
375 ///
376 /// The minimum size of the buffer can be found via [`Decoder::required_buf_len`].
377 #[inline]
378 pub fn decode_to_buf(&mut self, mut buf: impl AsMut<[u8]>) -> Result<usize> {
379 let buf = buf.as_mut();
380 let size = self.required_buf_len();
381 if unlikely(buf.len() < size) {
382 return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size });
383 }
384 self.reader.decode_image(buf, self.channels.as_u8(), self.header.channels.as_u8())?;
385 Ok(size)
386 }
387
388 /// Decodes the image into a newly allocated vector of bytes and returns it.
389 #[cfg(any(feature = "std", feature = "alloc"))]
390 #[inline]
391 pub fn decode_to_vec(&mut self) -> Result<Vec<u8>> {
392 let mut out = vec![0; self.header.n_pixels() * self.channels.as_u8() as usize];
393 let _ = self.decode_to_buf(&mut out)?;
394 Ok(out)
395 }
396}
397