Initial working version (but lots of temp stuff)

This commit is contained in:
Ivan Smirnov 2021-11-28 12:36:12 +00:00
commit f8ec8011fc
11 changed files with 753 additions and 0 deletions

3
.gitignore vendored Normal file
View file

@ -0,0 +1,3 @@
/target
Cargo.lock
/.idea

32
Cargo.toml Normal file
View file

@ -0,0 +1,32 @@
[package]
name = "qoi-fast"
version = "0.1.0"
authors = ["Ivan Smirnov <rust@ivan.smirnov.ie>"]
edition = "2021"
readme = "README.md"
repository = "https://github.com/aldanor/qoi-rust"
license = "MIT"
keywords = []
categories = []
description = "Pure Rust implementation of QOI (Quite Okay Image) format."
documentation = "https://docs.rs/qoi-rust"
[dev-dependencies]
png = "0.17"
[features]
reference-encoder = []
default = []
[target.'cfg(bench)'.dev-dependencies]
# to activate, pass RUSTFLAGS="--cfg bench" until cargo does this automatically
criterion = "^0.3.5"
[lib]
name = "qoi_fast"
path = "src/lib.rs"
bench = false
[[bench]]
name = "bench"
harness = false

58
benches/bench.rs Normal file
View file

@ -0,0 +1,58 @@
//! To run benchmarks, also pass RUSTFLAGS="--cfg bench" until cargo does this automatically.
use std::fs::File;
use criterion::{black_box, criterion_group, criterion_main, Criterion};
use qoi_fast::*;
pub fn criterion_benchmark(c: &mut Criterion) {
// let three_raw = include_bytes!("../assets/three.raw");
// let four_raw = include_bytes!("../assets/four.raw");
// let three_qoi = qoi_encode::<3>(three_raw, 572, 354).unwrap();
// c.bench_function("encode 3", |b| {
// b.iter(|| black_box(qoi_encode::<3>(three_raw, 572, 354).unwrap()))
// });
// c.bench_function("decode 3", |b| {
// b.iter(|| black_box(qoi_decode::<3>(&three_qoi).unwrap()))
// });
// c.bench_function("encode 4", |b| {
// b.iter(|| black_box(qoi_encode::<4>(four_raw, 572, 354).unwrap()))
// });
let decoder = png::Decoder::new(
File::open("/Users/ivansmirnov/projects/rust/qoi-other/images/kodak/kodim11.png").unwrap(),
);
let mut reader = decoder.read_info().unwrap();
let mut buf = vec![0; reader.output_buffer_size()];
let info = reader.next_frame(&mut buf).unwrap();
assert_eq!(info.color_type.samples(), 3);
let png_bytes = &buf[..info.buffer_size()];
let qoi_bytes = qoi_encode_to_vec(
png_bytes,
info.width as _,
info.height as _,
info.color_type.samples() as _,
)
.unwrap();
c.bench_function("kodim11.png (encode-3)", |b| {
b.iter(|| {
black_box(qoi_encode_to_vec(
png_bytes,
info.width as _,
info.height as _,
info.color_type.samples() as _,
))
.unwrap()
})
});
c.bench_function("kodim11.png (decode-3)", |b| {
b.iter(|| black_box(qoi_decode_to_vec(&qoi_bytes, 3)).unwrap())
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);

7
rustfmt.toml Normal file
View file

@ -0,0 +1,7 @@
use_small_heuristics = "Max"
use_field_init_shorthand = true
use_try_shorthand = true
empty_item_single_line = true
edition = "2018"
unstable_features = true
fn_args_layout = "Compressed"

16
src/consts.rs Normal file
View file

@ -0,0 +1,16 @@
pub const QOI_INDEX: u8 = 0x00;
pub const QOI_RUN_8: u8 = 0x40;
pub const QOI_RUN_16: u8 = 0x60;
pub const QOI_DIFF_8: u8 = 0x80;
pub const QOI_DIFF_16: u8 = 0xc0;
pub const QOI_DIFF_24: u8 = 0xe0;
pub const QOI_COLOR: u8 = 0xf0;
pub const QOI_MASK_2: u8 = 0xc0;
pub const QOI_MASK_3: u8 = 0xe0;
pub const QOI_MASK_4: u8 = 0xf0;
pub const QOI_HEADER_SIZE: usize = 14;
pub const QOI_PADDING: usize = 4;
pub const QOI_MAGIC: [u8; 4] = [b'q', b'o', b'i', b'f'];

157
src/decode.rs Normal file
View file

@ -0,0 +1,157 @@
use std::mem;
use crate::consts::{
QOI_COLOR, QOI_DIFF_16, QOI_DIFF_24, QOI_DIFF_8, QOI_HEADER_SIZE, QOI_INDEX, QOI_MAGIC,
QOI_MASK_2, QOI_MASK_3, QOI_MASK_4, QOI_PADDING, QOI_RUN_16, QOI_RUN_8,
};
use crate::error::{Error, Result};
use crate::header::Header;
use crate::pixel::{Pixel, SupportedChannels};
struct ReadBuf {
start: *const u8,
end: *const u8,
}
impl ReadBuf {
pub unsafe fn new(ptr: *const u8) -> Self {
Self { start: ptr, end: ptr }
}
#[inline]
pub fn read(&mut self) -> u8 {
unsafe {
let v = self.end.read();
self.end = self.end.add(1);
v
}
}
#[inline]
pub fn len(&self) -> usize {
unsafe { self.end.offset_from(self.start).max(0) as usize }
}
}
pub fn qoi_decode_impl<const N: usize>(data: &[u8]) -> Result<(Header, Vec<u8>)>
where
Pixel<N>: SupportedChannels,
{
if data.len() < QOI_HEADER_SIZE + QOI_PADDING {
return Err(Error::BadDecodingDataSize { size: data.len() });
}
let header = Header::from_bytes(unsafe {
// Safety: Header is a POD type and we have just checked that the data fits it
*(data.as_ptr() as *const _)
});
let n_pixels = (header.width as usize) * (header.height as usize);
if n_pixels == 0 {
return Err(Error::EmptyImage { width: header.width, height: header.height });
}
if header.magic != QOI_MAGIC {
return Err(Error::InvalidMagic { magic: header.magic });
}
let mut pixels = Vec::<Pixel<N>>::with_capacity(n_pixels);
unsafe {
// Safety: we have just allocated enough memory to set the length without problems
pixels.set_len(n_pixels);
}
let mut buf = unsafe {
// Safety: we will check within the loop that there are no reads outside the slice
ReadBuf::new(data.as_ptr().add(QOI_HEADER_SIZE))
};
let mut index = [Pixel::new(); 64];
let mut px = Pixel::new().with_a(0xff);
let mut run = 0_u16;
for px_out in pixels.iter_mut() {
// TODO: check for safety that ReadBuf is not over yet
if run != 0 {
run -= 1;
} else {
let b1 = buf.read();
match b1 & QOI_MASK_2 {
QOI_INDEX => {
px = index[usize::from(b1 ^ QOI_INDEX)];
}
QOI_DIFF_8 => {
px.rgb_add(
((b1 >> 4) & 0x03).wrapping_sub(2),
((b1 >> 2) & 0x03).wrapping_sub(2),
(b1 & 0x03).wrapping_sub(2),
);
}
_ => match b1 & QOI_MASK_3 {
QOI_RUN_8 => {
run = u16::from(b1 & 0x1f);
}
QOI_RUN_16 => {
run = 32 + ((u16::from(b1 & 0x1f) << 8) | u16::from(buf.read()));
}
QOI_DIFF_16 => {
let b2 = buf.read();
px.rgb_add(
(b1 & 0x1f).wrapping_sub(16),
(b2 >> 4).wrapping_sub(8),
(b2 & 0x0f).wrapping_sub(8),
);
}
_ => match b1 & QOI_MASK_4 {
QOI_DIFF_24 => {
let (b2, b3) = (buf.read(), buf.read());
px.rgba_add(
(((b1 & 0x0f) << 1) | (b2 >> 7)).wrapping_sub(16),
((b2 & 0x7c) >> 2).wrapping_sub(16),
(((b2 & 0x03) << 3) | ((b3 & 0xe0) >> 5)).wrapping_sub(16),
(b3 & 0x1f).wrapping_sub(16),
);
}
QOI_COLOR => {
if b1 & 8 != 0 {
px.set_r(buf.read());
}
if b1 & 4 != 0 {
px.set_g(buf.read());
}
if b1 & 2 != 0 {
px.set_b(buf.read());
}
if b1 & 1 != 0 {
px.set_a(buf.read());
}
}
_ => {}
},
},
}
unsafe {
// Safety: hash_index() is computed mod 64, so it will never go out of bounds
*index.get_unchecked_mut(usize::from(px.hash_index())) = px;
}
}
*px_out = px;
}
let bytes = unsafe {
// Safety: this is safe because we have previously set all the lengths ourselves
let ptr = pixels.as_mut_ptr();
mem::forget(pixels);
Vec::from_raw_parts(ptr as *mut _, n_pixels * N, n_pixels * N)
};
Ok((header, bytes))
}
pub fn qoi_decode_to_vec(data: impl AsRef<[u8]>, channels: u8) -> Result<(Header, Vec<u8>)> {
let data = data.as_ref();
match channels {
3 => qoi_decode_impl::<3>(data),
4 => qoi_decode_impl::<4>(data),
_ => Err(Error::InvalidChannels { channels }),
}
}

238
src/encode.rs Normal file
View file

@ -0,0 +1,238 @@
use std::slice;
use crate::consts::{
QOI_COLOR, QOI_DIFF_16, QOI_DIFF_24, QOI_DIFF_8, QOI_HEADER_SIZE, QOI_INDEX, QOI_PADDING,
QOI_RUN_16, QOI_RUN_8,
};
use crate::error::{Error, Result};
use crate::header::Header;
use crate::pixel::{Pixel, SupportedChannels};
struct WriteBuf {
start: *const u8,
current: *mut u8,
}
impl WriteBuf {
pub unsafe fn new(ptr: *mut u8) -> Self {
Self { start: ptr, current: ptr }
}
#[inline]
pub fn write<const N: usize>(&mut self, v: [u8; N]) {
unsafe {
let mut i = 0;
while i < N {
self.current.add(i).write(v[i]);
i += 1;
}
self.current = self.current.add(N);
}
}
#[inline]
pub fn push(&mut self, v: u8) {
unsafe {
self.current.write(v);
self.current = self.current.add(1);
}
}
#[inline]
pub fn len(&self) -> usize {
unsafe { self.current.offset_from(self.start).max(0) as usize }
}
}
#[inline]
fn encode_diff_reference<const N: usize>(
px: Pixel<N>, px_prev: Pixel<N>, buf: &mut WriteBuf,
) -> Option<(bool, bool, bool, bool)> {
let vr = (px.r() as i16) - (px_prev.r() as i16);
let vg = (px.g() as i16) - (px_prev.g() as i16);
let vb = (px.b() as i16) - (px_prev.b() as i16);
let va = (px.a_or(0) as i16) - (px_prev.a_or(0) as i16);
let (vr_16, vg_16, vb_16, va_16) = (vr + 16, vg + 16, vb + 16, va + 16);
if vr_16 | vg_16 | vb_16 | va_16 | 31 == 31 {
loop {
if va == 0 {
let (vr_2, vg_2, vb_2) = (vr + 2, vg + 2, vb + 2);
if vr_2 | vg_2 | vb_2 | 3 == 3 {
buf.write([QOI_DIFF_8 | (vr_2 << 4 | vg_2 << 2 | vb_2) as u8]);
break;
}
let (vg_8, vb_8) = (vg + 8, vb + 8);
if vg_8 | vb_8 | 15 == 15 {
buf.write([QOI_DIFF_16 | vr_16 as u8, (vg_8 << 4 | vb_8) as u8]);
break;
}
}
buf.write([
QOI_DIFF_24 | (vr_16 >> 1) as u8,
(vr_16 << 7 | vg_16 << 2 | vb_16 >> 3) as u8,
(vb_16 << 5 | va_16) as u8,
]);
break;
}
None
} else {
Some((vr != 0, vg != 0, vb != 0, va != 0))
}
}
#[inline]
fn encode_diff_wrapping<const N: usize>(
px: Pixel<N>, px_prev: Pixel<N>, buf: &mut WriteBuf,
) -> Option<(bool, bool, bool, bool)> {
let vr = px.r().wrapping_sub(px_prev.r());
let vg = px.g().wrapping_sub(px_prev.g());
let vb = px.b().wrapping_sub(px_prev.b());
let va = px.a_or(0).wrapping_sub(px_prev.a_or(0));
let (vr_16, vg_16, vb_16, va_16) =
(vr.wrapping_add(16), vg.wrapping_add(16), vb.wrapping_add(16), va.wrapping_add(16));
if vr_16 | vg_16 | vb_16 | va_16 | 31 == 31 {
loop {
if va == 0 {
let (vr_2, vg_2, vb_2) =
(vr.wrapping_add(2), vg.wrapping_add(2), vb.wrapping_add(2));
if vr_2 | vg_2 | vb_2 | 3 == 3 {
buf.write([QOI_DIFF_8 | vr_2 << 4 | vg_2 << 2 | vb_2]);
break;
}
let (vg_8, vb_8) = (vg.wrapping_add(8), vb.wrapping_add(8));
if vg_8 | vb_8 | 15 == 15 {
buf.write([QOI_DIFF_16 | vr_16, vg_8 << 4 | vb_8]);
break;
}
}
buf.write([
QOI_DIFF_24 | vr_16 >> 1,
vr_16 << 7 | vg_16 << 2 | vb_16 >> 3,
vb_16 << 5 | va_16,
]);
break;
}
None
} else {
Some((vr != 0, vg != 0, vb != 0, va != 0))
}
}
pub(crate) fn qoi_encode_impl<const N: usize>(
data: &[u8], width: u32, height: u32,
) -> Result<Vec<u8>>
where
Pixel<N>: SupportedChannels,
{
let n_pixels = (width as usize) * (height as usize);
if data.is_empty() {
return Err(Error::EmptyImage { width, height });
} else if n_pixels * N != data.len() {
return Err(Error::BadEncodingDataSize { size: data.len(), expected: n_pixels * N });
}
let pixels = unsafe {
// Safety: we've verified that n_pixels * N == data.len()
slice::from_raw_parts::<Pixel<N>>(data.as_ptr() as _, n_pixels)
};
let max_size = QOI_HEADER_SIZE + n_pixels * (N + 1) + QOI_PADDING;
let mut bytes = Vec::<u8>::with_capacity(max_size);
let mut buf = unsafe {
// Safety: all write ops are guaranteed to not go outside allocation
WriteBuf::new(bytes.as_mut_ptr())
};
let mut header = Header::default();
header.width = width;
header.height = height;
header.channels = N as u8;
// TODO: colorspace
buf.write(header.to_bytes());
let mut index = [Pixel::new(); 64];
let mut px_prev = Pixel::new().with_a(0xff);
let mut run = 0_u16;
let next_run = |buf: &mut WriteBuf, run: &mut u16| {
let mut r = *run;
if r < 33 {
r -= 1;
buf.push(QOI_RUN_8 | (r as u8));
} else {
r -= 33;
buf.write([QOI_RUN_16 | ((r >> 8) as u8), (r & 0xff) as u8]);
}
*run = 0;
};
for (i, &px) in pixels.iter().enumerate() {
if px == px_prev {
run += 1;
if run == 0x2020 || i == n_pixels - 1 {
next_run(&mut buf, &mut run);
}
} else {
if run != 0 {
next_run(&mut buf, &mut run);
}
let index_pos = px.hash_index();
let index_px = unsafe {
// Safety: hash_index() is computed mod 64, so it will never go out of bounds
index.get_unchecked_mut(usize::from(index_pos))
};
if *index_px == px {
buf.push(QOI_INDEX | index_pos);
} else {
*index_px = px;
let nonzero = if cfg!(feature = "reference-encoder") {
encode_diff_reference::<N>(px, px_prev, &mut buf)
} else {
encode_diff_wrapping::<N>(px, px_prev, &mut buf)
};
if let Some((r, g, b, a)) = nonzero {
let c = ((r as u8) << 3) | ((g as u8) << 2) | ((b as u8) << 1) | (a as u8);
buf.push(QOI_COLOR | c);
if r {
buf.push(px.r());
}
if g {
buf.push(px.g());
}
if b {
buf.push(px.b());
}
if a {
buf.push(px.a_or(0));
}
}
}
px_prev = px;
}
}
buf.write([0; QOI_PADDING]);
unsafe {
// Safety: buffer length cannot exceed allocated capacity
bytes.set_len(buf.len());
}
Ok(bytes)
}
pub fn qoi_encode_to_vec(
data: impl AsRef<[u8]>, width: u32, height: u32, channels: u8,
) -> Result<Vec<u8>> {
let data = data.as_ref();
match channels {
3 => qoi_encode_impl::<3>(data, width, height),
4 => qoi_encode_impl::<4>(data, width, height),
_ => Err(Error::InvalidChannels { channels }),
}
}

42
src/error.rs Normal file
View file

@ -0,0 +1,42 @@
use std::error::Error as StdError;
use std::fmt::{self, Display};
use std::result::Result as StdResult;
use crate::consts::{QOI_HEADER_SIZE, QOI_MAGIC, QOI_PADDING};
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Error {
InvalidChannels { channels: u8 },
EmptyImage { width: u32, height: u32 },
BadEncodingDataSize { size: usize, expected: usize },
BadDecodingDataSize { size: usize },
InvalidMagic { magic: [u8; 4] },
// TODO: invalid colorspace
}
pub type Result<T> = StdResult<T, Error>;
impl Display for Error {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self {
Self::InvalidChannels { channels } => {
write!(f, "invalid number of channels: {}", channels)
}
Self::EmptyImage { width, height } => {
write!(f, "image contains no pixels: {}x{}", width, height)
}
Self::BadEncodingDataSize { size, expected } => {
write!(f, "bad data size when encoding: {} (expected: {})", size, expected)
}
Self::BadDecodingDataSize { size } => {
let min_size = QOI_HEADER_SIZE + QOI_PADDING;
write!(f, "bad data size when decoding: {} (minimum required: {})", size, min_size)
}
Self::InvalidMagic { magic } => {
write!(f, "invalid magic: expected {:?}, got {:?}", QOI_MAGIC, magic)
}
}
}
}
impl StdError for Error {}

53
src/header.rs Normal file
View file

@ -0,0 +1,53 @@
use crate::consts::{QOI_HEADER_SIZE, QOI_MAGIC};
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct Header {
pub magic: [u8; 4],
pub width: u32,
pub height: u32,
pub channels: u8,
pub colorspace: u8,
}
impl Default for Header {
fn default() -> Self {
Self { magic: QOI_MAGIC, width: 0, height: 0, channels: 3, colorspace: 0 }
}
}
const fn u32_to_be(v: u32) -> [u8; 4] {
[
((0xff000000 & v) >> 24) as u8,
((0xff0000 & v) >> 16) as u8,
((0xff00 & v) >> 8) as u8,
(0xff & v) as u8,
]
}
const fn u32_from_be(v: &[u8]) -> u32 {
((v[0] as u32) << 24) | ((v[1] as u32) << 16) | ((v[2] as u32) << 8) | (v[3] as u32)
}
impl Header {
pub const SIZE: usize = QOI_HEADER_SIZE;
pub(crate) fn to_bytes(&self) -> [u8; QOI_HEADER_SIZE] {
let mut out = [0; QOI_HEADER_SIZE];
out[..4].copy_from_slice(&self.magic);
out[4..8].copy_from_slice(&u32_to_be(self.width));
out[8..12].copy_from_slice(&u32_to_be(self.height));
out[12] = self.channels;
out[13] = self.colorspace;
out
}
pub(crate) fn from_bytes(v: [u8; QOI_HEADER_SIZE]) -> Self {
let mut out = Self::default();
out.magic.copy_from_slice(&v[..4]);
out.width = u32_from_be(&v[4..8]);
out.height = u32_from_be(&v[8..12]);
out.channels = v[12];
out.colorspace = v[13];
out
}
}

58
src/lib.rs Normal file
View file

@ -0,0 +1,58 @@
mod consts;
mod decode;
mod encode;
mod error;
mod header;
mod pixel;
pub use crate::decode::qoi_decode_to_vec;
pub use crate::encode::qoi_encode_to_vec;
pub use crate::error::{Error, Result};
#[cfg(test)]
mod tests {
use std::fs::File;
use std::path::PathBuf;
use crate::{consts::QOI_MAGIC, qoi_decode_to_vec, qoi_encode_to_vec};
fn read_png(rel_path: &str) -> (u32, u32, u8, Vec<u8>) {
let get_path = || -> Option<PathBuf> {
Some(PathBuf::from(file!()).parent()?.parent()?.join(rel_path))
};
let decoder = png::Decoder::new(File::open(get_path().unwrap()).unwrap());
let mut reader = decoder.read_info().unwrap();
let mut buf = vec![0; reader.output_buffer_size()];
let info = reader.next_frame(&mut buf).unwrap();
let bytes = &buf[..info.buffer_size()];
(info.width, info.height, info.color_type.samples() as u8, bytes.to_vec())
}
#[test]
fn kodim_01() {
let (w, h, c, v) = read_png("assets/kodim01.png");
let q = qoi_encode_to_vec(&v, w, h, c).unwrap();
std::fs::write("kodim01.qoi", q.as_slice()).unwrap();
}
#[test]
fn wikipedia() {
let (w, h, c, v) = read_png("assets/en.wikipedia.org.png");
let q = qoi_encode_to_vec(&v, w, h, c).unwrap();
std::fs::write("wikipedia.qoi", q.as_slice()).unwrap();
}
#[test]
fn roundtrip_3() {
let three_raw = include_bytes!("../assets/three.raw").to_vec();
let (w, h, c) = (572, 354, 3);
let three_qoi = qoi_encode_to_vec(&three_raw, w, h, c).unwrap();
let (header, three_rtp) = qoi_decode_to_vec(&three_qoi, c).unwrap();
assert_eq!(header.magic, QOI_MAGIC);
assert_eq!(header.width, w);
assert_eq!(header.height, h);
assert_eq!(header.channels, c);
assert_eq!(three_rtp.len(), (w as usize) * (h as usize) * 3);
assert_eq!(three_raw, three_rtp.as_slice());
}
}

89
src/pixel.rs Normal file
View file

@ -0,0 +1,89 @@
#[derive(Copy, Clone, PartialEq, Eq)]
#[repr(transparent)]
pub struct Pixel<const N: usize>([u8; N]);
impl<const N: usize> Pixel<N> {
#[inline]
pub const fn new() -> Self {
Self([0; N])
}
#[inline]
pub const fn r(self) -> u8 {
self.0[0]
}
#[inline]
pub const fn g(self) -> u8 {
self.0[1]
}
#[inline]
pub const fn b(self) -> u8 {
self.0[2]
}
#[inline]
pub const fn with_a(mut self, value: u8) -> Self {
if N >= 4 {
self.0[3] = value;
}
self
}
#[inline]
pub const fn a_or(self, value: u8) -> u8 {
if N < 4 {
value
} else {
self.0[3]
}
}
#[inline]
pub const fn hash_index(self) -> u8 {
(self.r() ^ self.g() ^ self.b() ^ self.a_or(0xff)) % 64
}
#[inline]
pub fn rgb_add(&mut self, r: u8, g: u8, b: u8) {
self.0[0] = self.0[0].wrapping_add(r);
self.0[1] = self.0[1].wrapping_add(g);
self.0[2] = self.0[2].wrapping_add(b);
}
#[inline]
pub fn rgba_add(&mut self, r: u8, g: u8, b: u8, a: u8) {
self.rgb_add(r, g, b);
if N >= 4 {
self.0[3] = self.0[3].wrapping_add(a);
}
}
#[inline]
pub fn set_r(&mut self, value: u8) {
self.0[0] = value;
}
#[inline]
pub fn set_g(&mut self, value: u8) {
self.0[1] = value;
}
#[inline]
pub fn set_b(&mut self, value: u8) {
self.0[2] = value;
}
#[inline]
pub fn set_a(&mut self, value: u8) {
if N >= 4 {
self.0[3] = value;
}
}
}
pub trait SupportedChannels {}
impl SupportedChannels for Pixel<3> {}
impl SupportedChannels for Pixel<4> {}