pub fn encode(pixels: &[u32], width: usize, height: usize) -> Vec<u8> {
assert_eq!(pixels.len(), width * height);
let mut raw = Vec::with_capacity(height * (1 + width * 3));
for row in pixels.chunks(width) {
raw.push(0);
for &px in row {
raw.extend_from_slice(&[(px >> 16) as u8, (px >> 8) as u8, px as u8]);
}
}
let mut out = b"\x89PNG\r\n\x1a\n".to_vec();
let mut ihdr = Vec::with_capacity(13);
ihdr.extend_from_slice(&(width as u32).to_be_bytes());
ihdr.extend_from_slice(&(height as u32).to_be_bytes());
ihdr.extend_from_slice(&[8, 2, 0, 0, 0]);
chunk(&mut out, b"IHDR", &ihdr);
chunk(&mut out, b"IDAT", &zlib_stored(&raw));
chunk(&mut out, b"IEND", &[]);
out
}
fn chunk(out: &mut Vec<u8>, kind: &[u8; 4], data: &[u8]) {
out.extend_from_slice(&(data.len() as u32).to_be_bytes());
let start = out.len();
out.extend_from_slice(kind);
out.extend_from_slice(data);
let crc = crc32(&out[start..]);
out.extend_from_slice(&crc.to_be_bytes());
}
fn zlib_stored(data: &[u8]) -> Vec<u8> {
let mut out = vec![0x78, 0x01];
let mut blocks = data.chunks(65_535).peekable();
if blocks.peek().is_none() {
out.extend_from_slice(&[1, 0, 0, 0xFF, 0xFF]);
}
while let Some(block) = blocks.next() {
out.push(u8::from(blocks.peek().is_none()));
let len = block.len() as u16;
out.extend_from_slice(&len.to_le_bytes());
out.extend_from_slice(&(!len).to_le_bytes());
out.extend_from_slice(block);
}
out.extend_from_slice(&adler32(data).to_be_bytes());
out
}
fn crc32(data: &[u8]) -> u32 {
let mut crc = !0u32;
for &b in data {
crc ^= b as u32;
for _ in 0..8 {
crc = if crc & 1 != 0 { (crc >> 1) ^ 0xEDB8_8320 } else { crc >> 1 };
}
}
!crc
}
fn adler32(data: &[u8]) -> u32 {
let (mut a, mut b) = (1u32, 0u32);
for chunk in data.chunks(5552) {
for &x in chunk {
a += x as u32;
b += a;
}
a %= 65_521;
b %= 65_521;
}
(b << 16) | a
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn checksums() {
assert_eq!(crc32(b"IEND"), 0xAE42_6082);
assert_eq!(adler32(b"Wikipedia"), 0x11E6_0398);
}
#[test]
fn png_structure() {
let png = encode(&[0xFF0000, 0x00FF00, 0x0000FF, 0xFFFFFF], 2, 2);
assert!(png.starts_with(b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR\0\0\0\x02\0\0\0\x02\x08\x02"));
assert!(png.ends_with(b"IEND\xae\x42\x60\x82"));
let big = encode(&vec![0x123456; 640 * 400], 640, 400);
assert!(big.len() > 640 * 400 * 3);
}
}