use std::io::{Error, ErrorKind, Result, Write};
use crate::common::{hash, Channel, Colorspace, Rgb, Rgba};
/// Encodes a qoi image
///
/// This encodes an image writing it to a [`Write`r](std::io::Write).
///
/// If you're writing to a [`File`](std::fs::File) or similar you
/// probably want to wrap it in a [`BufWriter`](std::io::BufWriter)
/// for performance reasons.
#[derive(Clone, Debug)]
pub struct QoiWriter<W: Write>
{
writer: W,
width: u32,
height: u32,
channels: Channel,
colorspace: Colorspace,
previous_color: Rgba,
index: [Rgba; 64],
pixels_seen: u64,
current_run_length: u8,
finished: bool,
}
impl<W: Write> Drop for QoiWriter<W>
{
fn drop(&mut self)
{
// Errors sadly cannot be handled here.
drop(self.finish());
drop(self.flush());
}
}
impl<W: Write> QoiWriter<W>
{
/// Creates a new [`QoiWriter`]
///
/// This function creates a new [`QoiWriter`] from a
/// [`Write`r](std::io::Write).
///
/// The width, height, number of channels and the colorspace has
/// to be given so that the correct header can be written. The
/// last two of these don't change the encoding in any way (which
/// also means that theoretically even if formally there is no
/// alpha channel you could write alpha values. It should be
/// obvious that you *really* shouldn't do that since it's not
/// standard conform).
///
/// # Errors
/// An error is returned if the header couldn't be written.
pub fn new(
mut writer: W,
width: u32,
height: u32,
channels: Channel,
colorspace: Colorspace,
) -> Result<Self>
{
let mut header = Vec::with_capacity(14);
header.extend([b'q', b'o', b'i', b'f']);
header.extend(width.to_be_bytes());
header.extend(height.to_be_bytes());
match channels
{
Channel::Rgb => header.push(3),
Channel::Rgba => header.push(4),
}
match colorspace
{
Colorspace::Srgb => header.push(0),
Colorspace::Linear => header.push(1),
}
writer.write_all(&header)?;
Ok(Self {
writer,
width,
height,
channels,
colorspace,
previous_color: Rgba::new(0, 0, 0, 255),
index: [Rgba::new(0, 0, 0, 0); 64],
pixels_seen: 0,
current_run_length: 0,
finished: false,
})
}
/// Writes a single pixel (with alpha)
///
/// Writes a single pixel or buffers it if this is necessary for
/// compression. Due to this buffering you have to
/// [`flush()`](Self::flush), [`close()`](Self::close) or
/// [`drop()`](std::mem::drop) the image to have it written
/// immediately.
///
/// # Errors
/// A error is returned if already all pixels were written or if
/// the underlying writer returned one.
pub fn write_rgba(&mut self, color: Rgba) -> Result<()>
{
self.write_inner(color)?;
if self.width as u64 * self.height as u64 == self.pixels_seen
{
self.finish()?;
}
self.previous_color = color;
self.pixels_seen += 1;
self.index[hash(color)] = color;
Ok(())
}
/// Writes a single pixel (without alpha)
///
/// Writes a single pixel or buffers it if this is necessary for
/// compression. Due to this buffering you have to
/// [`flush()`](Self::flush), [`close()`](Self::close) or
/// [`drop()`](std::mem::drop) the image to have it written
/// immediately.
///
/// The alpha value is taken from the previous pixel (in case one
/// is needed).
///
/// # Errors
/// A error is returned if already all pixels were written or if
/// the underlying writer returned one.
pub fn write_rgb(&mut self, color: Rgb) -> Result<()>
{
let Rgb { r, g, b } = color;
let a = self.previous_color.a;
self.write_rgba(Rgba { r, g, b, a })
}
/// Flushes the image
///
/// Flushes the image and the underlying writer. This should
/// never be necessary since it's done automatically when
/// [`closed`](Self::close) or [`dropped`](std::mem::drop).
///
/// # Errors
/// A error is returned if the underlying writer returned an
/// error.
pub fn flush(&mut self) -> Result<()>
{
self.write_run_length()?;
self.writer.flush()
}
/// Closes the image
///
/// Closes the image and the underlying writer. This
/// automatically done on [`drop`](std::mem::drop) when it goes
/// out of scope.
///
/// # Errors
/// Returns an error if the underlying writer returned one.
pub fn close(mut self) -> Result<()>
{
self.finish()?;
self.flush()
}
/// Returns the dimensions of the image
///
/// Returns the dimensions (`(width, height)`) of the image.
#[must_use]
pub const fn dimensions(&self) -> (u32, u32)
{
(self.width, self.height)
}
/// Returns the width of the image
#[must_use]
pub const fn width(&self) -> u32
{
self.width
}
/// Returns the height of the image
#[must_use]
pub const fn height(&self) -> u32
{
self.height
}
/// Returns the number of channels of the image
///
/// Returns whether the alpha channel is used ([`Channel::Rgba`])
/// or not ([`Channel::Rgb`]).
///
/// Even if the alpha channel is disabled, a [`Rgba`] is decoded
/// instead of an [`Rgb`](crate::common::Rgb) (the alpha channel
/// is then always `255`).
///
/// The result of this function changes nothing in the encoding.
#[must_use]
pub const fn channels(&self) -> Channel
{
self.channels
}
// The "sRGB" in the second paragraph shouldn't be in backticks.
#[allow(clippy::doc_markdown)]
/// Returns the colorspace of the image
///
/// Return whether all channels are linear
/// ([`Colorspace::Linear`]) or if it's sRGB and only the alpha
/// channel (if it exists) is linear ([`Colorspace::Srgb`]).
///
/// The result of this function changes nothing in the encoding.
#[must_use]
pub const fn colorspace(&self) -> Colorspace
{
self.colorspace
}
fn finish(&mut self) -> Result<()>
{
if self.finished
{
Ok(())
}
else
{
self.finished = true;
self.writer.write_all(&[0, 0, 0, 0, 0, 0, 0, 1])
}
}
// This is for allowing "dr_dg" and "db_dg". This is done due to
// them being the official names from the standard.
#[allow(clippy::similar_names)]
fn write_inner(&mut self, color: Rgba) -> Result<()>
{
if self.width as u64 * self.height as u64 <= self.pixels_seen
{
// TODO: Change error code to StorageFull once #86442 gets stabilized.
return Err(Error::new(ErrorKind::Other, "all pixels already written"));
}
if color.a == self.previous_color.a
{
// QOI_OP_RUN
if color == self.previous_color
{
self.current_run_length += 1;
if self.current_run_length == 62
{
self.write_run_length()?;
}
return Ok(());
}
// QOI_OP_INDEX
if let Some(index) = self.index.iter().position(|&x| x == color)
{
return self.writer.write_all(&[index as u8]);
}
// QOI_OP_DIFF
let dg = color.g.wrapping_sub(self.previous_color.g).wrapping_add(2);
let dr = color.r.wrapping_sub(self.previous_color.r).wrapping_add(2);
let db = color.b.wrapping_sub(self.previous_color.b).wrapping_add(2);
if dr < 4 && dg < 4 && db < 4
{
return self.writer.write_all(&[64 + dr * 16 + dg * 4 + db]);
}
// QOI_OP_LUMA
let dr_dg = dr.wrapping_sub(dg).wrapping_add(8);
let db_dg = db.wrapping_sub(dg).wrapping_add(8);
let dg = dg.wrapping_add(30);
if dg < 64 && dr_dg < 16 && db_dg < 16
{
return self.writer.write_all(&[128 + dg, dr_dg * 16 + db_dg]);
}
// QOI_OP_RGB
self.writer.write_all(&[0xfe, color.r, color.g, color.b])
}
else
{
// QOI_OP_RGBA
self.writer
.write_all(&[0xff, color.r, color.g, color.b, color.a])
}
}
fn write_run_length(&mut self) -> Result<()>
{
if self.current_run_length == 0
{
return Ok(());
}
let val = 3 * 64 + self.current_run_length - 1;
self.writer.write_all(&[val])
}
}