use std::{
io::{Cursor, Seek, SeekFrom, Write},
mem::MaybeUninit,
};
pub enum Channels {
Three,
Four,
}
impl Channels {
fn len(&self) -> usize {
match self {
Self::Three => 3,
Self::Four => 4,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq)]
struct Pixel {
r: u8,
g: u8,
b: u8,
a: u8,
}
impl Pixel {
fn new(r: u8, g: u8, b: u8, a: u8) -> Self {
Self { r, g, b, a }
}
fn hash(&self) -> usize {
(self.r ^ self.g ^ self.b ^ self.a) as usize
}
}
pub struct Qoi;
impl Qoi {
const PADDING: usize = 4;
const INDEX: u8 = 0;
const RUN_8: u8 = 0b0100_0000;
const RUN_16: u8 = 0b0110_0000;
const DIFF_8: u8 = 0b1000_0000;
const DIFF_16: u8 = 0b1100_0000;
const DIFF_24: u8 = 0b1110_0000;
const COLOR: u8 = 0b1111_0000;
const MASK_2: u8 = 0b1100_0000;
const MASK_3: u8 = 0b1110_0000;
const MASK_4: u8 = 0b1111_0000;
}
#[repr(C)]
struct QoiHeader {
magic: [u8; 4],
width: [u8; 2],
height: [u8; 2],
size: [u8; 4],
}
impl QoiHeader {
const SIZE: usize = std::mem::size_of::<QoiHeader>();
fn new(width: usize, height: usize, size: usize) -> Self {
Self {
magic: *b"qoif",
width: u16::try_from(width).unwrap().to_be_bytes(),
height: u16::try_from(height).unwrap().to_be_bytes(),
size: u32::try_from(size).unwrap().to_be_bytes(),
}
}
fn as_slice(&self) -> &[u8] {
unsafe { std::slice::from_raw_parts(self as *const Self as *const u8, Self::SIZE) }
}
fn width(&self) -> u16 {
u16::from_be_bytes(self.width)
}
fn height(&self) -> u16 {
u16::from_be_bytes(self.height)
}
fn size(&self) -> u32 {
u32::from_be_bytes(self.size)
}
}
impl From<&[u8; QoiHeader::SIZE]> for QoiHeader {
fn from(input: &[u8; QoiHeader::SIZE]) -> Self {
let mut header = MaybeUninit::<QoiHeader>::uninit();
let header = unsafe {
std::ptr::copy(
input.as_ptr(),
header.as_mut_ptr() as *mut u8,
QoiHeader::SIZE,
);
header.assume_init()
};
assert_eq!(&input[0..4], b"qoif");
header
}
}
trait Between: PartialOrd
where
Self: Sized,
{
#[inline]
fn between(&self, low: Self, high: Self) -> bool {
*self >= low && *self <= high
}
}
impl Between for i16 {}
pub trait QoiEncode {
fn qoi_encode(
&self,
width: usize,
height: usize,
channels: Channels,
dest: impl AsMut<[u8]>,
) -> std::io::Result<()>;
}
impl<S> QoiEncode for S
where
S: AsRef<[u8]>,
{
fn qoi_encode(
&self,
width: usize,
height: usize,
channels: Channels,
mut dest: impl AsMut<[u8]>,
) -> std::io::Result<()> {
let max_size = (width * height) * (channels.len() + 1) + QoiHeader::SIZE + Qoi::PADDING;
let mut cursor = Cursor::new(dest.as_mut());
cursor.seek(SeekFrom::Start(QoiHeader::SIZE as u64))?;
let mut cache = [Pixel::default(); 64];
let mut previous_pixel = Pixel::default();
let mut run = 0u16;
for chunk in self.as_ref().chunks_exact(channels.len() as usize) {
let a = if channels.len() == 4 { chunk[3] } else { 0 };
let pixel = Pixel::new(chunk[0], chunk[1], chunk[2], a);
if pixel == previous_pixel {
run += 1;
}
if run > 0 && (run == 0x2020 || pixel != previous_pixel) {
if run < 33 {
run -= 1;
cursor.write_all(&[Qoi::RUN_8 | (run as u8)])?;
} else {
run -= 33;
cursor.write_all(&[Qoi::RUN_16 | ((run >> 8u16) as u8), run as u8])?;
}
run = 0;
}
if pixel != previous_pixel {
let cache_index = pixel.hash();
let cached_pixel = &mut cache[cache_index];
if &pixel == cached_pixel {
cursor
.write_all(&[Qoi::INDEX | (cache_index as u8)])
.unwrap();
} else {
*cached_pixel = pixel;
let dr = (pixel.r - previous_pixel.r) as i16;
let dr8 = dr as u8;
let dg = (pixel.g - previous_pixel.g) as i16;
let dg8 = dg as u8;
let db = (pixel.b - previous_pixel.b) as i16;
let db8 = db as u8;
let da = (pixel.a - previous_pixel.a) as i16;
let da8 = da as u8;
if da == 0 && dr.between(-1, 2) && dg.between(-1, 2) && db.between(-1, 2) {
cursor.write_all(&[Qoi::DIFF_8
| ((dr8 + 1) << 4)
| ((dg8 + 1) << 2)
| (db8 + 1)])?;
} else if da == 0
&& dr.between(-15, 16)
&& dg.between(-7, 8)
&& db.between(-7, 8)
{
cursor.write_all(&[
Qoi::DIFF_16 | (dr8 + 15),
((dg8 + 7) << 4) | (db8 + 7),
])?;
} else if dr.between(-15, 16)
&& dg.between(-15, 16)
&& db.between(-15, 16)
&& da.between(-15, 16)
{
cursor.write_all(&[
Qoi::DIFF_24 | ((dr8 + 15) >> 1),
((dr8 + 15) << 7) | ((dg8 + 15) << 2) | ((db8 + 15) >> 3),
((db8 + 15) << 5) | (da8 + 15),
])?;
} else {
let command = Qoi::COLOR
| if dr != 0 { 8 } else { 0 }
| if dg != 0 { 4 } else { 0 }
| if db != 0 { 2 } else { 0 }
| if da != 0 { 1 } else { 0 };
cursor.write_all(&[command, pixel.r, pixel.g, pixel.b, pixel.a])?;
}
}
}
previous_pixel = pixel;
}
cursor.write_all(&[0, 0, 0, 0])?;
let header = QoiHeader::new(width, height, cursor.position() as usize - QoiHeader::SIZE);
cursor.write_all(header.as_slice())?;
Ok(())
}
}
pub trait QoiDecode {
fn qoi_decode(
&self,
channels: Channels,
dest: impl AsMut<[u8]>,
) -> std::io::Result<(usize, usize)>;
}
impl<S> QoiDecode for S
where
S: AsRef<[u8]>,
{
fn qoi_decode(
&self,
channels: Channels,
dest: impl AsMut<[u8]>,
) -> std::io::Result<(usize, usize)> {
todo!()
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn round_trip_four_channels() {
const INPUT: &[u8] = include_bytes!("../swirl.qoi");
let mut decoded = [0u8; INPUT.len()];
INPUT
.qoi_encode(800, 800, Channels::Four, &mut decoded)
.unwrap();
let mut encoded = [0u8; INPUT.len()];
decoded.qoi_decode(Channels::Three, &mut encoded).unwrap();
assert_eq!(INPUT, decoded);
}
}