266 lines
11 KiB
Rust
266 lines
11 KiB
Rust
//! A PNG writer for raster exports: 8-bit RGBA (RGB when every pixel is opaque) or indexed
|
||
//! (PNG-8: a palette of up to 256 colours, 1–8 bits per pixel), adaptive row filters, the
|
||
//! resolution as a `pHYs` chunk and optional Adam7 interlacing (the image builds up progressively
|
||
//! while it loads). Chunks are written here; flate2 compresses the image data.
|
||
|
||
use std::borrow::Cow;
|
||
use std::io::Write as _;
|
||
|
||
use flate2::Crc;
|
||
use flate2::write::ZlibEncoder;
|
||
|
||
use super::quantize::Indexed;
|
||
|
||
/// What goes into the file besides the pixels.
|
||
#[derive(Clone, Copy, Debug, Default, PartialEq)]
|
||
pub struct PngOptions {
|
||
/// Pixels per inch, stored as pixels per metre in `pHYs` (`None`: no `pHYs` chunk).
|
||
pub ppi: Option<f64>,
|
||
/// Adam7 interlacing.
|
||
pub interlaced: bool,
|
||
}
|
||
|
||
const SIGNATURE: &[u8; 8] = b"\x89PNG\r\n\x1a\n";
|
||
/// Most image data bytes per `IDAT` chunk.
|
||
const IDAT_MAX: usize = 1 << 20;
|
||
/// Adam7 passes: (x0, y0, dx, dy).
|
||
const ADAM7: [(u32, u32, u32, u32); 7] = [(0, 0, 8, 8), (4, 0, 8, 8), (0, 4, 4, 8), (2, 0, 4, 4), (0, 2, 2, 4), (1, 0, 2, 2), (0, 1, 1, 2)];
|
||
|
||
/// Pixels per metre for a resolution in pixels per inch (the unit `pHYs` stores).
|
||
pub fn pixels_per_metre(ppi: f64) -> u32 {
|
||
(ppi / 0.0254).round().clamp(1.0, u32::MAX as f64) as u32
|
||
}
|
||
|
||
/// A whole `pHYs` chunk (length, type, data, CRC) declaring `x` × `y` pixels per inch.
|
||
pub fn phys_chunk(x: f64, y: f64) -> Vec<u8> {
|
||
let [x, y] = [x, y].map(|v| pixels_per_metre(v).to_be_bytes());
|
||
let mut out = Vec::with_capacity(21);
|
||
// x, y, unit 1 = metre.
|
||
chunk(&mut out, b"pHYs", &[x.as_slice(), y.as_slice(), &[1]].concat());
|
||
out
|
||
}
|
||
|
||
/// Encode straight-alpha RGBA8 pixels (`width`×`height`, row-major) as a PNG file.
|
||
pub fn encode(rgba: &[u8], width: u32, height: u32, o: &PngOptions) -> Result<Vec<u8>, String> {
|
||
check_size(width, height)?;
|
||
if rgba.len() as u64 != u64::from(width) * u64::from(height) * 4 {
|
||
return Err("PNG encoding failed: pixel buffer doesn't match the image size".into());
|
||
}
|
||
let px = rgba.as_chunks::<4>().0;
|
||
let opaque = px.iter().all(|p| p[3] == 255);
|
||
let (channels, color_type) = if opaque { (3, 2) } else { (4, 6) };
|
||
let pixels: Cow<[u8]> = if opaque { Cow::Owned(px.iter().flat_map(|p| [p[0], p[1], p[2]]).collect()) } else { Cow::Borrowed(rgba) };
|
||
write(&pixels, width, height, Layout { channels, bits: 8, color_type }, &[], o)
|
||
}
|
||
|
||
/// Encode a palette image as an indexed PNG (colour type 3: `PLTE`, and `tRNS` for its
|
||
/// transparent entry) at 1, 2, 4 or 8 bits per pixel, as few as the palette needs.
|
||
pub fn encode_indexed(img: &Indexed, o: &PngOptions) -> Result<Vec<u8>, String> {
|
||
check_size(img.width, img.height)?;
|
||
if img.indices.len() as u64 != u64::from(img.width) * u64::from(img.height) || img.palette.is_empty() || img.palette.len() > 256 {
|
||
return Err("PNG encoding failed: the palette image is malformed".into());
|
||
}
|
||
let bits = match img.palette.len() {
|
||
0..=2 => 1,
|
||
3..=4 => 2,
|
||
5..=16 => 4,
|
||
_ => 8,
|
||
};
|
||
let plte = img.palette.concat();
|
||
// Alpha of the entries up to the transparent one (the rest are opaque).
|
||
let trns: Vec<u8> = img.transparent.map(|t| (0..=t).map(|i| if i == t { 0 } else { 255 }).collect()).unwrap_or_default();
|
||
let mut extra: Vec<(&[u8; 4], &[u8])> = vec![(b"PLTE", &plte)];
|
||
if !trns.is_empty() {
|
||
extra.push((b"tRNS", &trns));
|
||
}
|
||
write(&img.indices, img.width, img.height, Layout { channels: 1, bits, color_type: 3 }, &extra, o)
|
||
}
|
||
|
||
fn check_size(width: u32, height: u32) -> Result<(), String> {
|
||
if width == 0 || height == 0 || width > i32::MAX as u32 || height > i32::MAX as u32 {
|
||
return Err(format!("PNG encoding failed: {width} × {height} pixels is not a valid image size"));
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// How pixels are stored: bytes per pixel (one per index for palette images), bits per sample
|
||
/// in the file and the PNG colour type.
|
||
#[derive(Clone, Copy)]
|
||
struct Layout {
|
||
channels: usize,
|
||
bits: u8,
|
||
color_type: u8,
|
||
}
|
||
|
||
/// Write the file: signature, `IHDR`, `pHYs`, the `extra` chunks, the filtered and compressed
|
||
/// pixels (`layout.channels` bytes each, packed to `layout.bits` in the file) and `IEND`.
|
||
fn write(pixels: &[u8], width: u32, height: u32, layout: Layout, extra: &[(&[u8; 4], &[u8])], o: &PngOptions) -> Result<Vec<u8>, String> {
|
||
let Layout { channels, bits, color_type } = layout;
|
||
let mut out = Vec::with_capacity(pixels.len() / 2 + 64);
|
||
out.extend_from_slice(SIGNATURE);
|
||
let mut ihdr = Vec::with_capacity(13);
|
||
ihdr.extend_from_slice(&width.to_be_bytes());
|
||
ihdr.extend_from_slice(&height.to_be_bytes());
|
||
// Bit depth, colour type, deflate compression, adaptive filtering, interlace method.
|
||
ihdr.extend_from_slice(&[bits, color_type, 0, 0, u8::from(o.interlaced)]);
|
||
chunk(&mut out, b"IHDR", &ihdr);
|
||
if let Some(ppi) = o.ppi {
|
||
out.extend_from_slice(&phys_chunk(ppi, ppi));
|
||
}
|
||
for (ty, data) in extra {
|
||
chunk(&mut out, ty, data);
|
||
}
|
||
// Filters work on whole bytes: a sub-byte pixel counts as one.
|
||
let bpp = (channels * bits as usize / 8).max(1);
|
||
let mut raw = Vec::with_capacity(pixels.len() + height as usize * if o.interlaced { 2 } else { 1 });
|
||
let mut rows = |px: &[u8], w: u32, h: u32| {
|
||
let packed = pack(px, w, h, bits);
|
||
filter_rows(&packed, (w as usize * channels * bits as usize).div_ceil(8), h, bpp, &mut raw);
|
||
};
|
||
if o.interlaced {
|
||
for pass in ADAM7 {
|
||
if let Some((w, h, px)) = sub_image(pixels, width, height, channels, pass) {
|
||
rows(&px, w, h);
|
||
}
|
||
}
|
||
} else {
|
||
rows(pixels, width, height);
|
||
}
|
||
let mut z = ZlibEncoder::new(Vec::with_capacity(raw.len() / 2), flate2::Compression::default());
|
||
let data = z.write_all(&raw).and_then(|()| z.finish()).map_err(|e| format!("PNG encoding failed: {e}"))?;
|
||
for part in data.chunks(IDAT_MAX) {
|
||
chunk(&mut out, b"IDAT", part);
|
||
}
|
||
chunk(&mut out, b"IEND", &[]);
|
||
Ok(out)
|
||
}
|
||
|
||
/// `png` (a PNG file) with a text chunk per `(keyword, text)` right after its header: `tEXt`, or
|
||
/// uncompressed `iTXt` (UTF-8) for text beyond Latin-1. Keywords are 1–79 printable ASCII
|
||
/// characters (others are skipped). Unchanged when it doesn't start with a PNG header.
|
||
pub fn with_text(png: Vec<u8>, entries: &[(&str, Cow<str>)]) -> Vec<u8> {
|
||
// The signature and the IHDR chunk (length, type, 13 data bytes, CRC), which comes first.
|
||
const HEADER: usize = 8 + 12 + 13;
|
||
if entries.is_empty() || !png.starts_with(SIGNATURE) || png.get(12..16) != Some(b"IHDR".as_slice()) || png.len() < HEADER {
|
||
return png;
|
||
}
|
||
let mut out = Vec::with_capacity(png.len() + entries.iter().map(|(k, v)| k.len() + v.len() + 20).sum::<usize>());
|
||
out.extend_from_slice(&png[..HEADER]);
|
||
for (keyword, text) in entries {
|
||
if keyword.is_empty() || keyword.len() > 79 || !keyword.bytes().all(|b| (32..=126).contains(&b)) {
|
||
continue;
|
||
}
|
||
let mut data = keyword.as_bytes().to_vec();
|
||
data.push(0);
|
||
let latin1: Option<Vec<u8>> = text.chars().map(|c| u8::try_from(u32::from(c)).ok()).collect();
|
||
let ty = match latin1 {
|
||
Some(bytes) => {
|
||
data.extend(bytes);
|
||
b"tEXt"
|
||
}
|
||
None => {
|
||
// Uncompressed, no language tag, no translated keyword.
|
||
data.extend_from_slice(&[0, 0, 0, 0]);
|
||
data.extend_from_slice(text.as_bytes());
|
||
b"iTXt"
|
||
}
|
||
};
|
||
if u32::try_from(data.len()).is_ok() {
|
||
chunk(&mut out, ty, &data);
|
||
}
|
||
}
|
||
out.extend_from_slice(&png[HEADER..]);
|
||
out
|
||
}
|
||
|
||
/// Rows of one-byte pixels packed `bits` per pixel, most significant first (each row starts on a
|
||
/// byte); at 8 bits the pixels as they are.
|
||
fn pack(px: &[u8], w: u32, h: u32, bits: u8) -> Cow<'_, [u8]> {
|
||
if bits >= 8 {
|
||
return Cow::Borrowed(px);
|
||
}
|
||
let per = (8 / bits) as usize;
|
||
let mut out = Vec::with_capacity(h as usize * (w as usize).div_ceil(per));
|
||
for row in px.chunks(w.max(1) as usize) {
|
||
for group in row.chunks(per) {
|
||
let byte = group.iter().enumerate().fold(0u8, |b, (i, v)| b | (v & ((1 << bits) - 1)) << (8 - bits as usize * (i + 1)));
|
||
out.push(byte);
|
||
}
|
||
}
|
||
Cow::Owned(out)
|
||
}
|
||
|
||
/// Append one chunk: length, type, data, CRC of type and data.
|
||
fn chunk(out: &mut Vec<u8>, ty: &[u8; 4], data: &[u8]) {
|
||
out.extend_from_slice(&(data.len() as u32).to_be_bytes());
|
||
out.extend_from_slice(ty);
|
||
out.extend_from_slice(data);
|
||
let mut crc = Crc::new();
|
||
crc.update(ty);
|
||
crc.update(data);
|
||
out.extend_from_slice(&crc.sum().to_be_bytes());
|
||
}
|
||
|
||
/// The pixels of one Adam7 pass as their own image (`None` when the pass is empty).
|
||
fn sub_image(px: &[u8], width: u32, height: u32, channels: usize, (x0, y0, dx, dy): (u32, u32, u32, u32)) -> Option<(u32, u32, Vec<u8>)> {
|
||
let w = width.saturating_sub(x0).div_ceil(dx);
|
||
let h = height.saturating_sub(y0).div_ceil(dy);
|
||
if w == 0 || h == 0 {
|
||
return None;
|
||
}
|
||
let mut out = Vec::with_capacity(w as usize * h as usize * channels);
|
||
for y in (y0..height).step_by(dy as usize) {
|
||
let row = &px[y as usize * width as usize * channels..][..width as usize * channels];
|
||
for x in (x0..width).step_by(dx as usize) {
|
||
out.extend_from_slice(&row[x as usize * channels..][..channels]);
|
||
}
|
||
}
|
||
Some((w, h, out))
|
||
}
|
||
|
||
/// Filter every row (`stride` bytes, `h` rows) into `out` (filter type byte + filtered bytes),
|
||
/// picking per row the filter with the smallest sum of absolute values (the usual heuristic).
|
||
fn filter_rows(px: &[u8], stride: usize, h: u32, bpp: usize, out: &mut Vec<u8>) {
|
||
let zero = vec![0u8; stride];
|
||
let mut cand: [Vec<u8>; 5] = std::array::from_fn(|_| vec![0u8; stride]);
|
||
for y in 0..h as usize {
|
||
let cur = &px[y * stride..][..stride];
|
||
let prev = if y == 0 { &zero[..] } else { &px[(y - 1) * stride..][..stride] };
|
||
// The first pixel has no left neighbour (a = c = 0).
|
||
let (head, tail) = (bpp.min(stride), stride.saturating_sub(bpp));
|
||
let [none, sub, up, avg, pae] = &mut cand;
|
||
none.copy_from_slice(cur);
|
||
sub[..head].copy_from_slice(&cur[..head]);
|
||
for i in 0..stride {
|
||
up[i] = cur[i].wrapping_sub(prev[i]);
|
||
}
|
||
for i in 0..head {
|
||
avg[i] = cur[i].wrapping_sub(prev[i] / 2);
|
||
pae[i] = cur[i].wrapping_sub(paeth(0, prev[i], 0));
|
||
}
|
||
for j in 0..tail {
|
||
let i = j + bpp;
|
||
let (a, b, c) = (cur[j], prev[i], prev[j]);
|
||
sub[i] = cur[i].wrapping_sub(a);
|
||
avg[i] = cur[i].wrapping_sub(((a as u16 + b as u16) / 2) as u8);
|
||
pae[i] = cur[i].wrapping_sub(paeth(a, b, c));
|
||
}
|
||
let cost = |c: &[u8]| c.iter().map(|&v| (v as i8).unsigned_abs() as u64).sum::<u64>();
|
||
let best = (0..5).min_by_key(|&f| cost(&cand[f])).unwrap_or(0);
|
||
out.push(best as u8);
|
||
out.extend_from_slice(&cand[best]);
|
||
}
|
||
}
|
||
|
||
fn paeth(a: u8, b: u8, c: u8) -> u8 {
|
||
let p = a as i16 + b as i16 - c as i16;
|
||
let (pa, pb, pc) = ((p - a as i16).abs(), (p - b as i16).abs(), (p - c as i16).abs());
|
||
if pa <= pb && pa <= pc {
|
||
a
|
||
} else if pb <= pc {
|
||
b
|
||
} else {
|
||
c
|
||
}
|
||
}
|