qoi/
decode.rs

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::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::new(&data)?;
130    let out = 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::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 = 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 = [0; QOI_HEADER_SIZE];
268        self.read_exact(&mut b)?;
269        Header::decode(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(Bytes::new(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}
qoi/decode.rs

qoi/
decode.rs
