#![allow(unsafe_code)]
use crate::huffman::ZIGZAG;
use crate::idct::dequant_idct_into;
use crate::{Jpeg, decode};
use pith_digest::fnv1a64;
use pith_math::idct2_2d;
pub const PITH_OK: i32 = 0;
pub const PITH_E_INVALID: i32 = -1;
pub const PITH_E_REJECTED: i32 = -2;
pub const PITH_LAYOUT_GRAY8: u32 = 0;
pub const PITH_LAYOUT_RGB8: u32 = 2;
pub const PITH_LAYOUT_RGBA8: u32 = 4;
const LANE: usize = 8;
const BLOCK: usize = 64;
const TWO_BLOCK: usize = 128;
unsafe fn read_lanes(buf: *const u8, n: usize) -> Vec<f64> {
let raw = unsafe { core::slice::from_raw_parts(buf, n * LANE) };
raw.chunks_exact(LANE)
.map(|c| f64::from_bits(u64::from_le_bytes(c.try_into().expect("8 bytes"))))
.collect()
}
unsafe fn write_lanes(buf: *mut u8, values: &[f64]) {
let raw = unsafe { core::slice::from_raw_parts_mut(buf, values.len() * LANE) };
for (dst, v) in raw.chunks_exact_mut(LANE).zip(values) {
dst.copy_from_slice(&v.to_bits().to_le_bytes());
}
}
fn lanes_in(buf: *const u8, len: usize, want: usize) -> Result<Vec<f64>, i32> {
if buf.is_null() || len != want * LANE {
return Err(PITH_E_INVALID);
}
Ok(unsafe { read_lanes(buf, want) })
}
unsafe fn lanes_out(out: *mut u8, out_len: usize, want: usize, values: &[f64]) -> Result<(), i32> {
if out.is_null() || out_len != want * LANE {
return Err(PITH_E_INVALID);
}
unsafe { write_lanes(out, values) };
Ok(())
}
fn digest_for(bytes: &[u8], layout: u32) -> Result<(u64, u32), i32> {
match layout {
PITH_LAYOUT_GRAY8 | PITH_LAYOUT_RGB8 | PITH_LAYOUT_RGBA8 => {}
_ => return Err(PITH_E_INVALID),
}
let jpeg = decode(bytes).map_err(|_| PITH_E_REJECTED)?;
let native_channels: u32 = match &jpeg {
Jpeg::Gray(_) => 1,
Jpeg::Rgb(_) => 3,
};
let bytes: alloc::vec::Vec<u8> = match (&jpeg, layout) {
(Jpeg::Gray(img), PITH_LAYOUT_GRAY8) => img.as_slice().to_vec(),
(Jpeg::Rgb(img), PITH_LAYOUT_RGB8) => img.as_slice().to_vec(),
(Jpeg::Gray(img), PITH_LAYOUT_RGB8) => {
img.as_slice().iter().flat_map(|&g| [g, g, g]).collect()
}
(Jpeg::Gray(img), PITH_LAYOUT_RGBA8) => img
.as_slice()
.iter()
.flat_map(|&g| [g, g, g, u8::MAX])
.collect(),
(Jpeg::Rgb(img), PITH_LAYOUT_RGBA8) => img
.as_slice()
.chunks_exact(3)
.flat_map(|c| [c[0], c[1], c[2], u8::MAX])
.collect(),
(Jpeg::Rgb(_), PITH_LAYOUT_GRAY8) => return Err(PITH_E_REJECTED),
_ => unreachable!("layout membership checked above"),
};
Ok((fnv1a64(&bytes), native_channels))
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pith_jpeg_decode(
data: *const u8,
len: usize,
layout: u32,
digest: *mut u64,
) -> i32 {
if data.is_null() && len != 0 {
return PITH_E_INVALID;
}
if digest.is_null() {
return PITH_E_INVALID;
}
if data.is_null() {
return PITH_E_REJECTED;
}
let bytes = unsafe { core::slice::from_raw_parts(data, len) };
match digest_for(bytes, layout) {
Ok((value, _)) => {
unsafe { *digest = value };
PITH_OK
}
Err(status) => status,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pith_jpeg_channels(data: *const u8, len: usize, out: *mut u32) -> i32 {
if data.is_null() && len != 0 {
return PITH_E_INVALID;
}
if out.is_null() {
return PITH_E_INVALID;
}
if data.is_null() {
return PITH_E_REJECTED;
}
let bytes = unsafe { core::slice::from_raw_parts(data, len) };
match decode(bytes) {
Ok(jpeg) => {
unsafe {
*out = match jpeg {
Jpeg::Gray(_) => 1,
Jpeg::Rgb(_) => 3,
};
}
PITH_OK
}
Err(_) => PITH_E_REJECTED,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pith_jpeg_dequant_zigzag(
input: *const u8,
input_len: usize,
output: *mut u8,
output_len: usize,
) -> i32 {
let payload = match lanes_in(input, input_len, BLOCK) {
Ok(v) => v,
Err(status) => return status,
};
let mut natural = alloc::vec![0.0; BLOCK];
for (k, &val) in payload.iter().enumerate() {
natural[ZIGZAG[k] as usize] = val;
}
match unsafe { lanes_out(output, output_len, BLOCK, &natural) } {
Ok(()) => PITH_OK,
Err(status) => status,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pith_jpeg_dequant(
input: *const u8,
input_len: usize,
output: *mut u8,
output_len: usize,
) -> i32 {
let lanes = match lanes_in(input, input_len, TWO_BLOCK) {
Ok(v) => v,
Err(status) => return status,
};
let products: Vec<f64> = lanes[..BLOCK]
.iter()
.zip(&lanes[BLOCK..])
.map(|(&c, &q)| c * q)
.collect();
match unsafe { lanes_out(output, output_len, BLOCK, &products) } {
Ok(()) => PITH_OK,
Err(status) => status,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pith_jpeg_idct_islow(
input: *const u8,
input_len: usize,
output: *mut u8,
output_len: usize,
) -> i32 {
let lanes = match lanes_in(input, input_len, TWO_BLOCK) {
Ok(v) => v,
Err(status) => return status,
};
let coefs: Vec<i32> = lanes[..BLOCK].iter().map(|&c| c as i32).collect();
let qt: Vec<u16> = lanes[BLOCK..].iter().map(|&q| q as u16).collect();
let mut plane = [0u8; BLOCK];
dequant_idct_into(&coefs, &qt, &mut plane, 8, 0, 0);
let samples: Vec<f64> = plane.iter().map(|&s| f64::from(s)).collect();
match unsafe { lanes_out(output, output_len, BLOCK, &samples) } {
Ok(()) => PITH_OK,
Err(status) => status,
}
}
#[unsafe(no_mangle)]
pub unsafe extern "C" fn pith_jpeg_idct_oracle(
input: *const u8,
input_len: usize,
output: *mut u8,
output_len: usize,
) -> i32 {
let mut block = match lanes_in(input, input_len, BLOCK) {
Ok(v) => v,
Err(status) => return status,
};
idct2_2d(&mut block, 8, 8);
match unsafe { lanes_out(output, output_len, BLOCK, &block) } {
Ok(()) => PITH_OK,
Err(status) => status,
}
}
#[cfg(test)]
mod tests {
use super::{
BLOCK, PITH_E_INVALID, PITH_E_REJECTED, PITH_LAYOUT_GRAY8, PITH_LAYOUT_RGB8,
PITH_LAYOUT_RGBA8, PITH_OK, digest_for, lanes_in, lanes_out, pith_jpeg_channels,
pith_jpeg_decode, pith_jpeg_dequant, pith_jpeg_dequant_zigzag, pith_jpeg_idct_islow,
pith_jpeg_idct_oracle,
};
use crate::reference;
fn lanes(values: &[f64]) -> Vec<u8> {
values
.iter()
.flat_map(|v| v.to_bits().to_le_bytes())
.collect()
}
fn unlanes(buf: &[u8]) -> Vec<f64> {
buf.chunks_exact(8)
.map(|c| f64::from_bits(u64::from_le_bytes(c.try_into().expect("8 bytes"))))
.collect()
}
fn fixture(name: &str) -> Vec<u8> {
std::fs::read(format!(
"{}/tests/fixtures/{name}",
env!("CARGO_MANIFEST_DIR")
))
.expect("fixture")
}
#[test]
fn ffi_digests_reproduce_every_recorded_value() {
for (name, want) in reference::digest_values() {
let file = reference::digests()
.iter()
.find(|d| d.name == name)
.expect("recorded digest")
.file;
let bytes = fixture(file);
let layout = if name.contains("gray") {
PITH_LAYOUT_GRAY8
} else {
PITH_LAYOUT_RGB8
};
let (digest, _) = digest_for(&bytes, layout).expect("decode");
assert_eq!(digest, want, "{name}");
}
}
#[test]
fn ffi_decode_statuses() {
let bytes = fixture("base_gray.jpg");
let mut digest: u64 = 0;
let status = unsafe {
pith_jpeg_decode(bytes.as_ptr(), bytes.len(), PITH_LAYOUT_GRAY8, &mut digest)
};
assert_eq!(status, PITH_OK);
let want = reference::digest_values()
.into_iter()
.find(|(n, _)| *n == "base_gray")
.expect("recorded")
.1;
assert_eq!(digest, want);
assert_eq!(want, 0xe0ce_77f0_0e2c_b066);
let mut channels: u32 = 0;
let status = unsafe { pith_jpeg_channels(bytes.as_ptr(), bytes.len(), &mut channels) };
assert_eq!(status, PITH_OK);
assert_eq!(channels, 1);
let garbage = b"not a jpeg at all";
let status = unsafe {
pith_jpeg_decode(
garbage.as_ptr(),
garbage.len(),
PITH_LAYOUT_GRAY8,
&mut digest,
)
};
assert_eq!(status, PITH_E_REJECTED);
let empty = [];
let status = unsafe { pith_jpeg_decode(empty.as_ptr(), 0, PITH_LAYOUT_GRAY8, &mut digest) };
assert_eq!(status, PITH_E_REJECTED);
let status = unsafe { pith_jpeg_decode(bytes.as_ptr(), bytes.len(), 9, &mut digest) };
assert_eq!(status, PITH_E_INVALID);
let color = fixture("base_444.jpg");
let status = unsafe {
pith_jpeg_decode(color.as_ptr(), color.len(), PITH_LAYOUT_GRAY8, &mut digest)
};
assert_eq!(status, PITH_E_REJECTED);
let status =
unsafe { pith_jpeg_decode(core::ptr::null(), 0, PITH_LAYOUT_GRAY8, &mut digest) };
assert_eq!(status, PITH_E_REJECTED);
let status = unsafe {
pith_jpeg_decode(
bytes.as_ptr(),
bytes.len(),
PITH_LAYOUT_GRAY8,
core::ptr::null_mut(),
)
};
assert_eq!(status, PITH_E_INVALID);
}
#[test]
fn ffi_layout_conversions() {
let gray = fixture("base_gray.jpg");
let mut native: u64 = 0;
let mut widened: u64 = 0;
unsafe {
assert_eq!(
pith_jpeg_decode(gray.as_ptr(), gray.len(), PITH_LAYOUT_GRAY8, &mut native),
PITH_OK
);
assert_eq!(
pith_jpeg_decode(gray.as_ptr(), gray.len(), PITH_LAYOUT_RGB8, &mut widened),
PITH_OK
);
assert_ne!(native, widened);
assert_eq!(
pith_jpeg_decode(gray.as_ptr(), gray.len(), PITH_LAYOUT_RGBA8, &mut widened),
PITH_OK
);
}
let color = fixture("base_444.jpg");
let mut rgba: u64 = 0;
unsafe {
assert_eq!(
pith_jpeg_decode(color.as_ptr(), color.len(), PITH_LAYOUT_RGBA8, &mut rgba),
PITH_OK
);
}
assert_ne!(rgba, 0);
}
#[test]
fn ffi_kernels_reproduce_the_recorded_vectors() {
for vector in reference::vectors() {
let input_lanes = lanes(&vector.input);
let mut output = vec![0u8; vector.output.len() * 8];
let status = unsafe {
match vector.name {
"dequant.zigzag.8x8" => pith_jpeg_dequant_zigzag(
input_lanes.as_ptr(),
input_lanes.len(),
output.as_mut_ptr(),
output.len(),
),
"dequant.8x8" => pith_jpeg_dequant(
input_lanes.as_ptr(),
input_lanes.len(),
output.as_mut_ptr(),
output.len(),
),
"idct.islow.8x8" => pith_jpeg_idct_islow(
input_lanes.as_ptr(),
input_lanes.len(),
output.as_mut_ptr(),
output.len(),
),
"idct.oracle.8x8" => pith_jpeg_idct_oracle(
input_lanes.as_ptr(),
input_lanes.len(),
output.as_mut_ptr(),
output.len(),
),
other => panic!("unmapped vector {other}"),
}
};
assert_eq!(status, PITH_OK, "{}", vector.name);
let got = unlanes(&output);
if vector.exact {
assert_eq!(got, vector.output, "{}", vector.name);
} else {
for (i, (&g, &w)) in got.iter().zip(&vector.output).enumerate() {
let budget = vector.tol_abs.max(vector.tol_rel * w.abs());
assert!(
(g - w).abs() <= budget,
"{} lane {i}: {g} vs {w}",
vector.name
);
}
}
}
}
#[test]
fn ffi_kernel_refusals() {
let ok_in = vec![0u8; BLOCK * 8];
let mut ok_out = vec![0u8; BLOCK * 8];
let short_in = vec![0u8; (BLOCK - 1) * 8];
let status = unsafe {
pith_jpeg_dequant_zigzag(
short_in.as_ptr(),
short_in.len(),
ok_out.as_mut_ptr(),
ok_out.len(),
)
};
assert_eq!(status, PITH_E_INVALID);
let status = unsafe {
pith_jpeg_dequant_zigzag(
ok_in.as_ptr(),
ok_in.len(),
ok_out.as_mut_ptr(),
ok_out.len() - 8,
)
};
assert_eq!(status, PITH_E_INVALID);
let status = unsafe {
pith_jpeg_dequant_zigzag(
core::ptr::null(),
ok_in.len(),
ok_out.as_mut_ptr(),
ok_out.len(),
)
};
assert_eq!(status, PITH_E_INVALID);
let status = unsafe {
pith_jpeg_dequant(
ok_in.as_ptr(),
ok_in.len(),
ok_out.as_mut_ptr(),
ok_out.len(),
)
};
assert_eq!(status, PITH_E_INVALID);
let two = vec![0u8; 128 * 8];
let status = unsafe {
pith_jpeg_idct_islow(two.as_ptr(), two.len(), ok_out.as_mut_ptr(), ok_out.len())
};
assert_eq!(status, PITH_OK);
let status = unsafe {
pith_jpeg_idct_oracle(
ok_in.as_ptr(),
ok_in.len(),
ok_out.as_mut_ptr(),
ok_out.len(),
)
};
assert_eq!(status, PITH_OK);
}
#[test]
fn lane_helpers_roundtrip() {
assert_eq!(
lanes_in(core::ptr::null(), 0, BLOCK).unwrap_err(),
PITH_E_INVALID
);
let values = vec![-0.5f64, 1.5, f64::MAX, -1.0e-300];
let bytes = lanes(&values);
let mut out = vec![0u8; bytes.len()];
unsafe { lanes_out(out.as_mut_ptr(), out.len(), values.len(), &values).expect("write") };
assert_eq!(unlanes(&out), values);
unsafe {
assert_eq!(
lanes_out(out.as_mut_ptr(), out.len() - 1, values.len(), &values).unwrap_err(),
PITH_E_INVALID
);
}
}
}