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>(buf:out, data),
82 Channels::Rgba => encode_impl::<_, 4>(buf: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: usize, height: usize) = (width as usize, height as usize);
92 let n_pixels: usize = 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}
211