Low budget PNG encoder

png.rs

#![allow(non_snake_case)]

use std::io::{Cursor, Error, Result, Write};

macro_rules! write_val {
	($writer:expr, $val:expr, $bigEndian:expr) => {{
		let bytes = if $bigEndian {
			$val.to_be_bytes()
		} else {
			$val.to_le_bytes()
		};
		$writer.write_all(&bytes)
	}};
	($writer:expr, $val:expr) => {
		write_val!($writer, $val, cfg!(target_endian = "big"))
	};
}

pub type ImgSize = (usize, usize);
type OutBuffer = Cursor<Vec<u8>>;

/// Create a 16-bit grayscale PNG in memory.
pub fn encode_grayscale_png(size: ImgSize, pixels: &[u16]) -> Result<Vec<u8>> {
	if size.0 >= u32::MAX as usize || size.1 >= u32::MAX as usize {
		return Err(Error::other("image dimensions too large"));
	}
	if size.0 * size.1 != pixels.len() {
		return Err(Error::other("pixel slice length does not match image size"));
	}

	let mut result = Cursor::new(Vec::with_capacity(1 << 10));
	write_header(&mut result, size, 16)?;
	write_idats(&mut result, size, unpacked_16bit_pixels(pixels), 2)?;
	Ok(result.into_inner())
}

fn unpacked_16bit_pixels(pixels: &[u16]) -> Vec<u8> {
	let mut result = Vec::with_capacity(2 * pixels.len());
	for &val in pixels {
		let bytes = val.to_be_bytes();
		result.write_all(&bytes).unwrap();
	}
	result
}

/// Create a full-color RGBA PNG in memory. Expects colors in (A)RGB order.
/// If alpha is false, still creates a png with alpha but sets all pixels to
/// maximum opacity.
pub fn encode_color_png(size: ImgSize, pixels: &[u32], alpha: bool) -> Result<Vec<u8>> {
	if size.0 >= u32::MAX as usize || size.1 >= u32::MAX as usize {
		return Err(Error::other("image dimensions too large"));
	}
	if size.0 * size.1 != pixels.len() {
		return Err(Error::other("pixel slice length does not match image size"));
	}

	let mut result = Cursor::new(Vec::with_capacity(1 << 10));
	write_header(&mut result, size, 32)?;
	write_idats(&mut result, size, unpacked_32bit_pixels(pixels, alpha), 4)?;
	Ok(result.into_inner())
}

fn unpacked_32bit_pixels(pixels: &[u32], alpha: bool) -> Vec<u8> {
	let mut result = Vec::with_capacity(4 * pixels.len());
	for &val in pixels {
		let mut bytes = val.to_be_bytes();

		if !alpha {
			// lazy: always write alpha but default to full opacity
			bytes[0] = 0xFF;
		}

		// png format expects RGBA, but we assume ARGB source, so we have to do some
		// swapping
		let bytes = [bytes[1], bytes[2], bytes[3], bytes[0]];
		result.write_all(&bytes).unwrap();
	}
	result
}

/// Helper to dump an encoded png onto disk.
pub fn write_to_file(encoded: &Vec<u8>, path: &str) -> Result<()> {
	let mut f = std::fs::File::create(path)?;
	f.write_all(encoded.as_slice())?;
	Ok(())
}

fn write_chunk(buf: &mut OutBuffer, sig: &[u8], data: &[u8]) -> Result<()> {
	assert_eq!(sig.len(), 4, "chunk signature is 4 bytes");
	let mut crc = crc32(sig);
	crc = crc32_partial(crc, data);

	write_val!(buf, data.len() as u32, true)?;
	buf.write_all(sig)?;
	buf.write_all(data)?;
	write_val!(buf, crc, true)?;

	Ok(())
}

fn write_header(buf: &mut OutBuffer, size: ImgSize, bitDepth: u8) -> Result<()> {
	assert!(bitDepth == 16 || bitDepth == 32, "unknown bit depth");

	let mut subBuf = Cursor::new(Vec::with_capacity(1 << 10));
	write_val!(subBuf, size.0 as u32, true)?; // width
	write_val!(subBuf, size.1 as u32, true)?; // height
	write_val!(subBuf, if bitDepth == 32 { 8 } else { 16u8 })?; // bit depth
	write_val!(subBuf, if bitDepth == 32 { 6 } else { 0u8 })?; // color type
	write_val!(subBuf, 0u8)?; // compression method
	write_val!(subBuf, 0u8)?; // filter method
	write_val!(subBuf, 0u8)?; // interlace method

	buf.write_all(b"\x89PNG\r\n\x1A\n")?; // file signature
	write_chunk(buf, b"IHDR", &subBuf.into_inner())?;
	Ok(())
}

fn write_idats(
	buf: &mut OutBuffer,
	size: ImgSize,
	pixels: Vec<u8>,
	bytesPerPixel: usize,
) -> Result<()> {
	let filtered = filter_scanlines(pixels, size, bytesPerPixel);
	let zlibCrc = adler32(filtered.as_slice()); // checks _uncompressed_ data
	let mut deflated = dumb_deflate(filtered);
	// write zlib header
	deflated.insert(0, 0x78);
	deflated.insert(1, 0x01);
	// write zlib trailer
	write_val!(deflated, zlibCrc, true)?;

	// image data may need to be spread among several IDATs
	let mut slice = deflated.as_slice();
	while slice.len() > 0 {
		let subslice = &slice[0 .. usize::min(slice.len(), 0xFFFFFF)];
		slice = &slice[subslice.len() ..];
		write_chunk(buf, b"IDAT", subslice)?;
	}
	write_chunk(buf, b"IEND", &[])?;

	Ok(())
}

// Prepends each scanline with a filter-type byte indicating no filtering.
fn filter_scanlines(mut buf: Vec<u8>, size: ImgSize, bytesPerPixel: usize) -> Vec<u8> {
	for y in (0 .. size.1).rev() {
		let idx = y * size.0 * bytesPerPixel;
		buf.insert(idx, 0); // filter type 0
	}
	buf
}

// Passes data through uncompressed with bare-minimum DEFLATE headers.
fn dumb_deflate(buf: Vec<u8>) -> Vec<u8> {
	let mut result = Vec::with_capacity(buf.len() + 5 * (buf.len() / 0xFFFF)); // copying of huge buffers :sad_pickle:
	let mut slice = buf.as_slice();

	while slice.len() > 0 {
		let sliceLen = usize::min(slice.len(), 0xFFFF);
		let subslice = &slice[.. sliceLen];
		slice = &slice[sliceLen ..];

		result
			.write_all(&[
				// DEFLATE block continuation bit
				if slice.len() > 0 { 0 } else { 1 },
			])
			.unwrap();
		write_val!(result, subslice.len() as u16, false).unwrap(); // no idea why these are little-endian unlike everything else
		write_val!(result, !(subslice.len() as u16), false).unwrap(); // the spec even seems to indicate otherwise?
		result.write_all(subslice).unwrap();
	}

	result
}

use self::hash::*;
mod hash {
	use std::sync::LazyLock;

	/// Compute Adler32 checksum of the given bytes.
	/// Because we really needed two separate hashing algorithms in our image
	/// format
	pub fn adler32(buf: &[u8]) -> u32 {
		const MOD_ADLER: u32 = 65521;
		
		let mut a = 1;
		let mut b = 0;
		for &byte in buf {
			a = (a + byte as u32) % MOD_ADLER;
			b = (b + a) % MOD_ADLER;
		}
		(b << 16) | a
	}

	/// Compute CRC32 checksum of the given bytes.
	pub fn crc32(buf: &[u8]) -> u32 {
		crc32_partial(0, buf)
	}

	/// Compute a partial CRC32 over the block of bytes, given the partial
	/// checksum from prior blocks.
	pub fn crc32_partial(last: u32, buf: &[u8]) -> u32 {
		static CRC_TABLE: LazyLock<[u32; 256]> = LazyLock::new(|| {
			let mut table = [0; 256];
			let mut val;
			for index in 0 .. 256 {
				val = index as u32;
				for _ in 0 .. 8 {
					if val & 1 > 0 {
						val = 0xEDB88320 ^ ((val >> 1) & 0x7FFFFFFF);
					} else {
						val = (val >> 1) & 0x7FFFFFFF;
					}
				}
				table[index] = val;
			}
			table
		});

		let mut crc = last ^ 0xFFFFFFFF;
		for &byte in buf {
			crc = CRC_TABLE[((crc ^ byte as u32) & 0xFF) as usize] ^ (crc >> 8);
		}
		crc ^ 0xFFFFFFFF
	}

	#[cfg(test)]
	mod test {
		use super::*;

		#[test]
		fn test_adler() {
			let buf: &[u8] = "foobarbaz".as_bytes();

			assert_eq!(adler32(&[]), 1);
			assert_eq!(adler32(&buf[0 .. 3]), 42074437);
			assert_eq!(adler32(&buf[3 .. 6]), 39649590);
			assert_eq!(adler32(&buf[6 .. 9]), 40173886);
			assert_eq!(adler32(buf), 310051767);
		}

		#[test]
		fn test_crc() {
			let buf: &[u8] = "foobarbaz".as_bytes();

			assert_eq!(crc32(&[]), 0);
			assert_eq!(crc32(&buf[0 .. 3]), 2356372769);
			assert_eq!(crc32(&buf[3 .. 6]), 1996459178);
			assert_eq!(crc32(&buf[6 .. 9]), 2015626392);
			assert_eq!(crc32(buf), 444082090);

			// running checksums
			assert_eq!(crc32_partial(crc32(&buf[0 .. 3]), &buf[3 .. 6]), 2666930069);
			assert_eq!(
				crc32_partial(
					crc32_partial(crc32(&buf[0 .. 3]), &buf[3 .. 6]),
					&buf[6 .. 9]
				),
				444082090
			);
		}
	}
}