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

qoi/
decode.rs
