use std::convert::TryFrom; use std::io::Write; use crate::consts::{QOI_HEADER_SIZE, QOI_OP_INDEX, QOI_OP_RUN, QOI_PADDING, QOI_PADDING_SIZE}; use crate::error::{Error, Result}; use crate::header::Header; use crate::pixel::{Pixel, SupportedChannels}; use crate::types::{Channels, ColorSpace}; use crate::utils::{unlikely, BytesMut, GenericWriter, Writer}; #[allow(clippy::cast_possible_truncation)] fn qoi_encode_impl(mut buf: W, data: &[u8]) -> Result where Pixel: SupportedChannels, { let cap = buf.capacity(); let mut index = [Pixel::new(); 256]; let mut px_prev = Pixel::new().with_a(0xff); let mut run = 0_u8; let mut px = Pixel::::new().with_a(0xff); let n_pixels = data.len() / N; for (i, chunk) in data.chunks_exact(N).enumerate() { px.read(chunk); if px == px_prev { run += 1; if run == 62 || unlikely(i == n_pixels - 1) { buf = buf.write_one(QOI_OP_RUN | (run - 1))?; run = 0; } } else { if run != 0 { buf = buf.write_one(QOI_OP_RUN | (run - 1))?; run = 0; } let index_pos = px.hash_index(); let index_px = &mut index[index_pos as usize]; let px_rgba = px.as_rgba(0xff); if *index_px == px_rgba { buf = buf.write_one(QOI_OP_INDEX | index_pos)?; } else { *index_px = px_rgba; buf = px.encode_into(px_prev, buf)?; } px_prev = px; } } buf = buf.write_many(&QOI_PADDING)?; Ok(cap.saturating_sub(buf.capacity())) } #[inline] fn qoi_encode_impl_all(out: W, data: &[u8], channels: Channels) -> Result { match channels { Channels::Rgb => qoi_encode_impl::<_, 3>(out, data), Channels::Rgba => qoi_encode_impl::<_, 4>(out, data), } } #[inline] pub fn encoded_size_limit(width: u32, height: u32, channels: impl Into) -> usize { let (width, height) = (width as usize, height as usize); let n_pixels = width.saturating_mul(height); QOI_HEADER_SIZE + n_pixels.saturating_mul(channels.into() as usize) + n_pixels + QOI_PADDING_SIZE } #[inline] pub fn qoi_encode_to_buf( buf: impl AsMut<[u8]>, data: impl AsRef<[u8]>, width: u32, height: u32, ) -> Result { QoiEncoder::new(&data, width, height)?.encode_to_buf(buf) } #[inline] pub fn qoi_encode_to_vec(data: impl AsRef<[u8]>, width: u32, height: u32) -> Result> { QoiEncoder::new(&data, width, height)?.encode_to_vec() } pub struct QoiEncoder<'a> { data: &'a [u8], header: Header, } impl<'a> QoiEncoder<'a> { #[inline] #[allow(clippy::cast_possible_truncation)] pub fn new(data: &'a (impl AsRef<[u8]> + ?Sized), width: u32, height: u32) -> Result { let data = data.as_ref(); let mut header = Header::try_new(width, height, Channels::default(), ColorSpace::default())?; let size = data.len(); let n_channels = size / header.n_pixels(); if header.n_pixels() * n_channels != size { return Err(Error::InvalidImageLength { size, width, height }); } header.channels = Channels::try_from(n_channels.min(0xff) as u8)?; Ok(Self { data, header }) } #[inline] pub const fn with_colorspace(mut self, colorspace: ColorSpace) -> Self { self.header = self.header.with_colorspace(colorspace); self } #[inline] pub const fn channels(&self) -> Channels { self.header.channels } #[inline] pub const fn header(&self) -> &Header { &self.header } #[inline] pub fn encoded_size_limit(&self) -> usize { self.header.encoded_size_limit() } #[inline] pub fn encode_to_buf(&self, mut buf: impl AsMut<[u8]>) -> Result { let buf = buf.as_mut(); let size_required = self.encoded_size_limit(); if unlikely(buf.len() < size_required) { return Err(Error::OutputBufferTooSmall { size: buf.len(), required: size_required }); } let (head, tail) = buf.split_at_mut(QOI_HEADER_SIZE); // can't panic head.copy_from_slice(&self.header.encode()); let n_written = qoi_encode_impl_all(BytesMut::new(tail), self.data, self.header.channels)?; Ok(QOI_HEADER_SIZE + n_written) } #[inline] pub fn encode_to_vec(&self) -> Result> { let mut out = vec![0_u8; self.encoded_size_limit()]; let size = self.encode_to_buf(&mut out)?; out.truncate(size); Ok(out) } #[inline] pub fn encode_to_stream(&self, writer: &mut W) -> Result { writer.write_all(&self.header.encode())?; let n_written = qoi_encode_impl_all(GenericWriter::new(writer), self.data, self.header.channels)?; Ok(n_written + QOI_HEADER_SIZE) } }