qoi/
encode.rs

1#[cfg(any(feature = "std", feature = "alloc"))]
2use alloc::{vec, vec::Vec};
3use core::convert::TryFrom;
4#[cfg(feature = "std")]
5use std::io::Write;
6
7use bytemuck::Pod;
8
9use crate::consts::{QOI_HEADER_SIZE, QOI_OP_INDEX, QOI_OP_RUN, QOI_PADDING, QOI_PADDING_SIZE};
10use crate::error::{Error, Result};
11use crate::header::Header;
12use crate::pixel::{Pixel, SupportedChannels};
13use crate::types::{Channels, ColorSpace};
14#[cfg(feature = "std")]
15use crate::utils::GenericWriter;
16use crate::utils::{unlikely, BytesMut, Writer};
17
18#[allow(clippy::cast_possible_truncation, unused_assignments, unused_variables)]
19fn encode_impl<W: Writer, const N: usize>(mut buf: W, data: &[u8]) -> Result<usize>
20where
21    Pixel<N>: SupportedChannels,
22    [u8; N]: Pod,
23{
24    let cap = buf.capacity();
25
26    let mut index = [Pixel::new(); 256];
27    let mut px_prev = Pixel::new().with_a(0xff);
28    let mut hash_prev = px_prev.hash_index();
29    let mut run = 0_u8;
30    let mut px = Pixel::<N>::new().with_a(0xff);
31    let mut index_allowed = false;
32
33    let n_pixels = data.len() / N;
34
35    for (i, chunk) in data.chunks_exact(N).enumerate() {
36        px.read(chunk);
37        if px == px_prev {
38            run += 1;
39            if run == 62 || unlikely(i == n_pixels - 1) {
40                buf = buf.write_one(QOI_OP_RUN | (run - 1))?;
41                run = 0;
42            }
43        } else {
44            if run != 0 {
45                #[cfg(not(feature = "reference"))]
46                {
47                    // credits for the original idea: @zakarumych (had to be fixed though)
48                    buf = buf.write_one(if run == 1 && index_allowed {
49                        QOI_OP_INDEX | hash_prev
50                    } else {
51                        QOI_OP_RUN | (run - 1)
52                    })?;
53                }
54                #[cfg(feature = "reference")]
55                {
56                    buf = buf.write_one(QOI_OP_RUN | (run - 1))?;
57                }
58                run = 0;
59            }
60            index_allowed = true;
61            let px_rgba = px.as_rgba(0xff);
62            hash_prev = px_rgba.hash_index();
63            let index_px = &mut index[hash_prev as usize];
64            if *index_px == px_rgba {
65                buf = buf.write_one(QOI_OP_INDEX | hash_prev)?;
66            } else {
67                *index_px = px_rgba;
68                buf = px.encode_into(px_prev, buf)?;
69            }
70            px_prev = px;
71        }
72    }
73
74    buf = buf.write_many(&QOI_PADDING)?;
75    Ok(cap.saturating_sub(buf.capacity()))
76}
77
78#[inline]
79fn encode_impl_all<W: Writer>(out: W, data: &[u8], channels: Channels) -> Result<usize> {
80    match channels {
81        Channels::Rgb => encode_impl::<_, 3>(out, data),
82        Channels::Rgba => encode_impl::<_, 4>(out, data),
83    }
84}
85
86/// The maximum number of bytes the encoded image will take.
87///
88/// Can be used to pre-allocate the buffer to encode the image into.
89#[inline]
90pub fn encode_max_len(width: u32, height: u32, channels: impl Into<u8>) -> usize {
91    let (width, height) = (width as usize, height as usize);
92    let n_pixels = width.saturating_mul(height);
93    QOI_HEADER_SIZE
94        + n_pixels.saturating_mul(channels.into() as usize)
95        + n_pixels
96        + QOI_PADDING_SIZE
97}
98
99/// Encode the image into a pre-allocated buffer.
100///
101/// Returns the total number of bytes written.
102#[inline]
103pub fn encode_to_buf(
104    buf: impl AsMut<[u8]>, data: impl AsRef<[u8]>, width: u32, height: u32,
105) -> Result<usize> {
106    Encoder::new(&data, width, height)?.encode_to_buf(buf)
107}
108
109/// Encode the image into a newly allocated vector.
110#[cfg(any(feature = "alloc", feature = "std"))]
111#[inline]
112pub fn encode_to_vec(data: impl AsRef<[u8]>, width: u32, height: u32) -> Result<Vec<u8>> {
113    Encoder::new(&data, width, height)?.encode_to_vec()
114}
115
116/// Encode QOI images into buffers or into streams.
117pub struct Encoder<'a> {
118    data: &'a [u8],
119    header: Header,
120}
121
122impl<'a> Encoder<'a> {
123    /// Creates a new encoder from a given array of pixel data and image dimensions.
124    ///
125    /// The number of channels will be inferred automatically (the valid values
126    /// are 3 or 4). The color space will be set to sRGB by default.
127    #[inline]
128    #[allow(clippy::cast_possible_truncation)]
129    pub fn new(data: &'a (impl AsRef<[u8]> + ?Sized), width: u32, height: u32) -> Result<Self> {
130        let data = data.as_ref();
131        let mut header =
132            Header::try_new(width, height, Channels::default(), ColorSpace::default())?;
133        let size = data.len();
134        let n_channels = size / header.n_pixels();
135        if header.n_pixels() * n_channels != size {
136            return Err(Error::InvalidImageLength { size, width, height });
137        }
138        header.channels = Channels::try_from(n_channels.min(0xff) as u8)?;
139        Ok(Self { data, header })
140    }
141
142    /// Returns a new encoder with modified color space.
143    ///
144    /// Note: the color space doesn't affect encoding or decoding in any way, it's
145    /// a purely informative field that's stored in the image header.
146    #[inline]
147    pub const fn with_colorspace(mut self, colorspace: ColorSpace) -> Self {
148        self.header = self.header.with_colorspace(colorspace);
149        self
150    }
151
152    /// Returns the inferred number of channels.
153    #[inline]
154    pub const fn channels(&self) -> Channels {
155        self.header.channels
156    }
157
158    /// Returns the header that will be stored in the encoded image.
159    #[inline]
160    pub const fn header(&self) -> &Header {
161        &self.header
162    }
163
164    /// The maximum number of bytes the encoded image will take.
165    ///
166    /// Can be used to pre-allocate the buffer to encode the image into.
167    #[inline]
168    pub fn required_buf_len(&self) -> usize {
169        self.header.encode_max_len()
170    }
171
172    /// Encodes the image to a pre-allocated buffer and returns the number of bytes written.
173    ///
174    /// The minimum size of the buffer can be found via [`Encoder::required_buf_len`].
175    #[inline]
176    pub fn encode_to_buf(&self, mut buf: impl AsMut<[u8]>) -> Result<usize> {
177        let buf = buf.as_mut();
178        let size_required = self.required_buf_len();
179        if unlikely(buf.len() < size_required) {
180            return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size_required });
181        }
182        let (head, tail) = buf.split_at_mut(QOI_HEADER_SIZE); // can't panic
183        head.copy_from_slice(&self.header.encode());
184        let n_written = encode_impl_all(BytesMut::new(tail), self.data, self.header.channels)?;
185        Ok(QOI_HEADER_SIZE + n_written)
186    }
187
188    /// Encodes the image into a newly allocated vector of bytes and returns it.
189    #[cfg(any(feature = "alloc", feature = "std"))]
190    #[inline]
191    pub fn encode_to_vec(&self) -> Result<Vec<u8>> {
192        let mut out = vec![0_u8; self.required_buf_len()];
193        let size = self.encode_to_buf(&mut out)?;
194        out.truncate(size);
195        Ok(out)
196    }
197
198    /// Encodes the image directly to a generic writer that implements [`Write`](std::io::Write).
199    ///
200    /// Note: while it's possible to pass a `&mut [u8]` slice here since it implements `Write`,
201    /// it would more effficient to use a specialized method instead: [`Encoder::encode_to_buf`].
202    #[cfg(feature = "std")]
203    #[inline]
204    pub fn encode_to_stream<W: Write>(&self, writer: &mut W) -> Result<usize> {
205        writer.write_all(&self.header.encode())?;
206        let n_written =
207            encode_impl_all(GenericWriter::new(writer), self.data, self.header.channels)?;
208        Ok(n_written + QOI_HEADER_SIZE)
209    }
210}
qoi/encode.rs

qoi/
encode.rs
