From fb62a41730ba5862b2342a521656bd6eea91b879 Mon Sep 17 00:00:00 2001 From: Will Toohey Date: Sat, 25 Apr 2026 23:57:52 +1000 Subject: [PATCH] (vibe coded) move PNG work to rust --- Cargo.lock | 76 +++++++++++++++++++++ Cargo.toml | 3 +- pyproject.toml | 2 +- rust/lib.rs | 21 +++++- rust/png_enc.rs | 90 +++++++++++++++++++++++++ src/ifstools/handlers/image_decoders.py | 19 ++++++ src/ifstools/handlers/image_file.py | 6 +- src/ifstools/handlers/lz77.py | 2 +- src/ifstools/handlers/tex_folder.py | 6 +- 9 files changed, 212 insertions(+), 13 deletions(-) create mode 100644 rust/png_enc.rs diff --git a/Cargo.lock b/Cargo.lock index 169e9ae..1b6090a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,52 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "bitflags" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "fdeflate" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c" +dependencies = [ + "simd-adler32", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "heck" version = "0.5.0" @@ -12,6 +58,7 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" name = "ifstools-native" version = "0.1.0" dependencies = [ + "png", "pyo3", ] @@ -21,12 +68,35 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "once_cell" version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "png" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61" +dependencies = [ + "bitflags", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -109,6 +179,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + [[package]] name = "syn" version = "2.0.117" diff --git a/Cargo.toml b/Cargo.toml index 633bd7c..2d573b7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,11 +4,12 @@ version = "0.1.0" edition = "2021" [lib] -name = "_lz77_native" +name = "_native" crate-type = ["cdylib", "rlib"] path = "rust/lib.rs" [dependencies] +png = "0.18.1" pyo3 = { version = "0.28", features = ["extension-module", "abi3-py310"] } [profile.release] diff --git a/pyproject.toml b/pyproject.toml index bce11aa..cf31d2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,5 +26,5 @@ ifstools = "ifstools:main" [tool.maturin] python-source = "src" -module-name = "ifstools.handlers._lz77_native" +module-name = "ifstools.handlers._native" features = ["pyo3/extension-module"] diff --git a/rust/lib.rs b/rust/lib.rs index e4e1aab..3eca0d0 100644 --- a/rust/lib.rs +++ b/rust/lib.rs @@ -3,6 +3,7 @@ use pyo3::prelude::*; use pyo3::types::PyBytes; mod lz77; +mod png_enc; #[pyfunction] #[pyo3(name = "decompress")] @@ -25,10 +26,26 @@ fn py_compress<'py>( Ok(PyBytes::new(py, &out)) } +#[pyfunction] +#[pyo3(name = "encode_png", signature = (width, height, pixels, color="rgba"))] +fn py_encode_png<'py>( + py: Python<'py>, + width: u32, + height: u32, + pixels: Vec, + color: &str, +) -> PyResult> { + let out = py + .detach(|| png_enc::encode(width, height, &pixels, color)) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(PyBytes::new(py, &out)) +} + #[pymodule] -#[pyo3(name = "_lz77_native")] -fn _lz77_native(m: &Bound<'_, PyModule>) -> PyResult<()> { +#[pyo3(name = "_native")] +fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_decompress, m)?)?; m.add_function(wrap_pyfunction!(py_compress, m)?)?; + m.add_function(wrap_pyfunction!(py_encode_png, m)?)?; Ok(()) } diff --git a/rust/png_enc.rs b/rust/png_enc.rs new file mode 100644 index 0000000..94a569d --- /dev/null +++ b/rust/png_enc.rs @@ -0,0 +1,90 @@ +//! Thin wrapper over the `png` crate for fast PNG encoding from raw pixel +//! buffers. Used in place of Pillow's PNG path during IFS extraction. + +use png::{BitDepth, ColorType, Encoder}; + +#[derive(Debug)] +pub enum EncodeError { + UnknownColorType(String), + SizeMismatch { expected: usize, got: usize }, + Png(png::EncodingError), +} + +impl std::fmt::Display for EncodeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + EncodeError::UnknownColorType(s) => write!(f, "unknown color type: {}", s), + EncodeError::SizeMismatch { expected, got } => write!( + f, + "pixel buffer size mismatch (expected {} bytes, got {})", + expected, got + ), + EncodeError::Png(e) => write!(f, "png error: {}", e), + } + } +} + +impl std::error::Error for EncodeError {} + +impl From for EncodeError { + fn from(e: png::EncodingError) -> Self { + EncodeError::Png(e) + } +} + +fn parse_color(s: &str) -> Result<(ColorType, usize), EncodeError> { + match s { + "rgba" | "RGBA" => Ok((ColorType::Rgba, 4)), + "rgb" | "RGB" => Ok((ColorType::Rgb, 3)), + "la" | "LA" => Ok((ColorType::GrayscaleAlpha, 2)), + "l" | "L" => Ok((ColorType::Grayscale, 1)), + other => Err(EncodeError::UnknownColorType(other.to_string())), + } +} + +pub fn encode( + width: u32, + height: u32, + pixels: &[u8], + color: &str, +) -> Result, EncodeError> { + let (ct, bpp) = parse_color(color)?; + let expected = (width as usize) * (height as usize) * bpp; + if pixels.len() != expected { + return Err(EncodeError::SizeMismatch { + expected, + got: pixels.len(), + }); + } + + // Reasonable starting capacity: assume PNG is no larger than raw pixels. + let mut out = Vec::with_capacity(pixels.len()); + { + let mut encoder = Encoder::new(&mut out, width, height); + encoder.set_color(ct); + encoder.set_depth(BitDepth::Eight); + encoder.set_compression(png::Compression::Balanced); + let mut writer = encoder.write_header()?; + writer.write_image_data(pixels)?; + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rgba_smoke() { + let pixels: Vec = (0..4 * 4 * 4).map(|i| i as u8).collect(); + let png = encode(4, 4, &pixels, "rgba").unwrap(); + // PNG signature. + assert_eq!(&png[..8], &[137, 80, 78, 71, 13, 10, 26, 10]); + } + + #[test] + fn size_mismatch_errors() { + let pixels = vec![0u8; 16]; + assert!(encode(4, 4, &pixels, "rgba").is_err()); + } +} diff --git a/src/ifstools/handlers/image_decoders.py b/src/ifstools/handlers/image_decoders.py index decbabf..2c28208 100644 --- a/src/ifstools/handlers/image_decoders.py +++ b/src/ifstools/handlers/image_decoders.py @@ -5,6 +5,25 @@ from struct import pack from PIL import Image from tqdm import tqdm +try: + from . import _native +except ImportError: + _native = None + +# PIL modes we can pass through directly; anything else is converted to RGBA. +_PNG_DIRECT_MODES = {'RGBA', 'RGB', 'LA', 'L'} + +def encode_png(im): + '''Encode a PIL Image as PNG bytes via the Rust png crate when available, + falling back to PIL's encoder otherwise.''' + if _native is None: + b = BytesIO() + im.save(b, format='PNG') + return b.getvalue() + if im.mode not in _PNG_DIRECT_MODES: + im = im.convert('RGBA') + return _native.encode_png(im.width, im.height, im.tobytes(), im.mode.lower()) + # header for a standard DDS with DXT5 compression and RGBA pixels # gap placed for image height/width insertion dxt_start = b'DDS |\x00\x00\x00\x07\x10\x00\x00' diff --git a/src/ifstools/handlers/image_file.py b/src/ifstools/handlers/image_file.py index f3a0f45..e606f2d 100644 --- a/src/ifstools/handlers/image_file.py +++ b/src/ifstools/handlers/image_file.py @@ -16,7 +16,7 @@ from tqdm import tqdm from .. import utils from . import lz77 from .generic_file import GenericFile -from .image_decoders import cachable_formats, image_formats +from .image_decoders import cachable_formats, encode_png, image_formats class ImageFile(GenericFile): @@ -82,12 +82,10 @@ class ImageFile(GenericFile): ) im = im.crop(dims) - b = BytesIO() if raw_pixels: return (im.width, im.height), im.tobytes() else: - im.save(b, format = 'PNG') - return b.getvalue() + return encode_png(im) def repack(self, manifest, data_blob, tqdm_progress, **kwargs): if tqdm_progress: diff --git a/src/ifstools/handlers/lz77.py b/src/ifstools/handlers/lz77.py index baa7851..076b564 100644 --- a/src/ifstools/handlers/lz77.py +++ b/src/ifstools/handlers/lz77.py @@ -1,5 +1,5 @@ try: - from ._lz77_native import compress, decompress + from ._native import compress, decompress except ImportError: print("WARNING: using native-python LZ77, operations will be slow") from ._lz77_py import compress, decompress diff --git a/src/ifstools/handlers/tex_folder.py b/src/ifstools/handlers/tex_folder.py index 75f3d5c..d4afff3 100644 --- a/src/ifstools/handlers/tex_folder.py +++ b/src/ifstools/handlers/tex_folder.py @@ -5,7 +5,7 @@ from PIL import Image, ImageDraw from tqdm import tqdm from .generic_file import GenericFile -from .image_decoders import cachable_formats +from .image_decoders import cachable_formats, encode_png from .image_file import ImageFile from .md5_folder import MD5Folder @@ -54,9 +54,7 @@ class ImageCanvas(GenericFile): draw.rectangle((size[0], size[2], size[1], size[3]), outline='red') del draw - b = BytesIO() - im.save(b, format = 'PNG') - return b.getvalue() + return encode_png(im) # since it's basically metadata, we ignore similarly to _cache def repack(self, manifest, data_blob, tqdm_progress, **kwargs):