diff --git a/src/encode.rs b/src/encode.rs index 12d3300..5b7bbf3 100644 --- a/src/encode.rs +++ b/src/encode.rs @@ -1,17 +1,18 @@ +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}; +use crate::utils::{unlikely, BytesMut, GenericWriter, Writer}; #[allow(clippy::cast_possible_truncation)] -fn qoi_encode_impl(out: &mut [u8], data: &[u8]) -> Result +fn qoi_encode_impl(mut buf: W, data: &[u8]) -> Result where Pixel: SupportedChannels, { - let out_size = out.len(); - let mut buf = BytesMut::new(out); + let cap = buf.capacity(); let mut index = [Pixel::new(); 256]; let mut px_prev = Pixel::new().with_a(0xff); @@ -25,36 +26,36 @@ where if px == px_prev { run += 1; if run == 62 || unlikely(i == n_pixels - 1) { - buf = buf.write_one(QOI_OP_RUN | (run - 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)); + 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); + buf = buf.write_one(QOI_OP_INDEX | index_pos)?; } else { *index_px = px_rgba; - buf = px.encode_into(px_prev, buf); + buf = px.encode_into(px_prev, buf)?; } px_prev = px; } } - buf = buf.write_many(&QOI_PADDING); - Ok(out_size.saturating_sub(buf.len())) + buf = buf.write_many(&QOI_PADDING)?; + Ok(cap.saturating_sub(buf.capacity())) } #[inline] -fn qoi_encode_impl_all(out: &mut [u8], data: &[u8], channels: Channels) -> Result { +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), + Channels::Rgb => qoi_encode_impl::<_, 3>(out, data), + Channels::Rgba => qoi_encode_impl::<_, 4>(out, data), } } @@ -130,7 +131,7 @@ impl<'a> QoiEncoder<'a> { } 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(tail, self.data, self.header.channels)?; + let n_written = qoi_encode_impl_all(BytesMut::new(tail), self.data, self.header.channels)?; Ok(QOI_HEADER_SIZE + n_written) } @@ -141,4 +142,12 @@ impl<'a> QoiEncoder<'a> { 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) + } } diff --git a/src/pixel.rs b/src/pixel.rs index 092d25c..0f8fe7b 100644 --- a/src/pixel.rs +++ b/src/pixel.rs @@ -1,4 +1,5 @@ use crate::consts::{QOI_OP_DIFF, QOI_OP_LUMA, QOI_OP_RGB, QOI_OP_RGBA}; +use crate::error::Result; use crate::utils::Writer; #[derive(Copy, Clone, PartialEq, Eq)] @@ -118,7 +119,7 @@ impl Pixel { } #[inline] - pub fn encode_into(&self, px_prev: Self, buf: W) -> W::Output { + pub fn encode_into(&self, px_prev: Self, buf: W) -> Result { if N == 3 || self.a_or(0) == px_prev.a_or(0) { let vg = self.g().wrapping_sub(px_prev.g()); let vg_32 = vg.wrapping_add(32); diff --git a/src/utils.rs b/src/utils.rs index 1ce45f0..d790d80 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,3 +1,7 @@ +use std::io::Write; + +use crate::error::Result; + #[inline(always)] #[cold] pub const fn cold() {} @@ -19,24 +23,53 @@ pub const fn unlikely(b: bool) -> bool { b } -pub trait Writer { - type Output; +pub trait Writer: Sized { + fn write_one(self, v: u8) -> Result; + fn write_many(self, v: &[u8]) -> Result; + fn capacity(&self) -> usize; +} - fn write_one(self, v: u8) -> Self::Output; - fn write_many(self, v: &[u8]) -> Self::Output; +pub struct GenericWriter { + writer: W, + n_written: usize, +} + +impl GenericWriter { + pub fn new(writer: W) -> Self { + Self { writer, n_written: 0 } + } +} + +impl Writer for GenericWriter { + fn write_one(mut self, v: u8) -> Result { + self.n_written += 1; + self.writer.write_all(&[v]).map(|_| self).map_err(Into::into) + } + + fn write_many(mut self, v: &[u8]) -> Result { + self.n_written += v.len(); + self.writer.write_all(v).map(|_| self).map_err(Into::into) + } + + fn capacity(&self) -> usize { + usize::MAX - self.n_written + } } impl<'a> Writer for BytesMut<'a> { - type Output = BytesMut<'a>; - #[inline] - fn write_one(self, v: u8) -> Self { - BytesMut::write_one(self, v) + fn write_one(self, v: u8) -> Result { + Ok(BytesMut::write_one(self, v)) } #[inline] - fn write_many(self, v: &[u8]) -> Self { - BytesMut::write_many(self, v) + fn write_many(self, v: &[u8]) -> Result { + Ok(BytesMut::write_many(self, v)) + } + + #[inline] + fn capacity(&self) -> usize { + self.0.len() } } @@ -64,9 +97,4 @@ impl<'a> BytesMut<'a> { head.copy_from_slice(v); Self(tail) } - - #[inline] - pub const fn len(&self) -> usize { - self.0.len() - } }