(vibe coded) move PNG work to rust

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"

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]

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"]

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<u8>,
+    color: &str,
+) -> PyResult<Bound<'py, PyBytes>> {
+    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(())
 }

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<png::EncodingError> 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<Vec<u8>, 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<u8> = (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());
+    }
+}

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'

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:

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

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):