From 392b8720cd25ba00bb1cc9d3a8b3c895cfa04259 Mon Sep 17 00:00:00 2001 From: "Dylan M. Taylor" Date: Wed, 11 Nov 2026 21:04:07 -0500 Subject: [PATCH] craft-psd: less copying and compression work when encoding planes (#1043) - PackBits scans each row once; RLE planes encode their rows on worker threads on native targets (rayon), keeping the exact packet boundaries and the fallback when RLE does not pay. - ZIP output grows in bounded blocks and is compacted to its encoded size, instead of reserving the raw size; native threads reuse a finished encoder's working buffers. - 16/32-bit prediction works on owned buffers in one pass. - Filter-effects planes and pattern blocks encode from borrowed rows and owned buffers without the intermediate copies. The encoded bytes are unchanged. Ported from storytold/photocraft#7797 (its crates/psd part), at a1ee8355a4651d379286cd13f92dabe4344e295f. Co-authored-by: Rene Leonhardt <65483435+reneleonhardt@users.noreply.github.com> --- Cargo.lock | 1 + crates/psd/Cargo.toml | 3 + crates/psd/src/compression.rs | 295 ++++++++++++++++++++++------ crates/psd/src/compression/tests.rs | 292 +++++++++++++++++++++++++++ crates/psd/src/filter_effects.rs | 67 ++++++- crates/psd/src/patterns.rs | 61 +++--- crates/psd/src/patterns/tests.rs | 121 ++++++++++++ 7 files changed, 743 insertions(+), 97 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3a66892d..b6b4afa9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1711,6 +1711,7 @@ dependencies = [ "encoding_rs", "flate2", "proptest", + "rayon", "serde", "thiserror 2.0.21", ] diff --git a/crates/psd/Cargo.toml b/crates/psd/Cargo.toml index d9d02aba..c9aa6dbe 100644 --- a/crates/psd/Cargo.toml +++ b/crates/psd/Cargo.toml @@ -21,6 +21,9 @@ flate2 = { version = "1.1.10", default-features = false, features = ["zlib-rs", encoding_rs = { workspace = true } serde = { workspace = true, optional = true } +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +rayon = { workspace = true } + [dev-dependencies] proptest = { workspace = true } craft-psd = { path = ".", features = ["testgen"] } diff --git a/crates/psd/src/compression.rs b/crates/psd/src/compression.rs index 7d05f547..d684525a 100644 --- a/crates/psd/src/compression.rs +++ b/crates/psd/src/compression.rs @@ -68,34 +68,90 @@ pub mod packbits { /// Encodes one row, appending to `out`. Never emits the `-128` no-op. pub fn encode(src: &[u8], out: &mut Vec) { + // Constant 16-bit samples with different bytes contain only literal packets. + if src.len() >= 128 + && let Some(&[a, b]) = src.first_chunk::<2>() + && a != b + { + let (blocks, tail) = src.as_chunks::<8>(); + if blocks.iter().all(|v| *v == [a, b, a, b, a, b, a, b]) && tail.iter().zip([a, b].into_iter().cycle()).all(|(&v, p)| v == p) { + for chunk in src.chunks(128) { + out.push((chunk.len() - 1) as u8); + out.extend_from_slice(chunk); + } + return; + } + } + packets(src, |header, bytes| { + out.push(header); + out.extend_from_slice(bytes); + }); + } + + /// The encoded row length, without allocating or writing its payload. + pub fn encoded_len(src: &[u8]) -> usize { + let mut len = 0; + packets(src, |_, bytes| len += 1 + bytes.len()); + len + } + + /// Visit PackBits run and literal packets in their deterministic encoding order. + fn packets(src: &[u8], mut emit: impl FnMut(u8, &[u8])) { let n = src.len(); let mut i = 0; while i < n { // Length of the run starting at i. let mut run = 1; - while i + run < n && run < 128 && src[i + run] == src[i] { + while i + run < n && run < 8 && src[i + run] == src[i] { run += 1; } + if run == 8 { + while run + 8 <= 128 && src.get(i + run..).and_then(|s| s.first_chunk::<8>()) == Some(&[src[i]; 8]) { + run += 8; + } + while i + run < n && run < 128 && src[i + run] == src[i] { + run += 1; + } + } // A literal never stops before a 2-byte run, so any run of >= 2 // seen here starts a packet boundary: a run packet is never larger. if run >= 2 { - out.push((1i16 - run as i16) as i8 as u8); - out.push(src[i]); + emit((1i16 - run as i16) as i8 as u8, &src[i..i + 1]); i += run; continue; } // Literal: extend until a run of >= 3 begins or 128 bytes. let start = i; let mut len = 0; - while i < n && len < 128 { + // Short packets keep the scalar scan; lookahead pays only on longer literals. + while i < n && len < 16 { if i + 2 < n && src[i] == src[i + 1] && src[i] == src[i + 2] { break; } i += 1; len += 1; } - out.push((len - 1) as u8); - out.extend_from_slice(&src[start..start + len]); + if len == 16 { + while len + 8 <= 128 { + let Some(&[a, b, c, d, e, f, g, h, j, k]) = src.get(i..).and_then(|s| s.first_chunk::<10>()) else { break }; + let word = u64::from_le_bytes([a, b, c, d, e, f, g, h]); + let different = (word ^ u64::from_le_bytes([b, c, d, e, f, g, h, j])) | (word ^ u64::from_le_bytes([c, d, e, f, g, h, j, k])); + // A zero byte marks a triple; scalar scanning keeps the first boundary. + if different.wrapping_sub(0x0101010101010101) & !different & 0x8080808080808080 != 0 { + break; + } + i += 8; + len += 8; + } + while i < n && len < 128 { + if i + 2 < n && src[i] == src[i + 1] && src[i] == src[i + 2] { + break; + } + i += 1; + len += 1; + } + } + emit((len - 1) as u8, &src[start..start + len]); } } @@ -369,82 +425,178 @@ fn decode_rle(data: &[u8], layout: &PlaneLayout, total: usize) -> Result Ok(out) } +/// Encode independent row chunks and join their count tables and payloads in order. fn encode_rle(decoded: &[u8], layout: &PlaneLayout) -> Result> { let rows = layout.rows()?; - let rb = layout.row_bytes(); - let cs = layout.count_size(); if rows == 0 { return Ok(Vec::new()); } - // Rows are independent, so they pack in parallel chunks of rows. Two passes: the first only - // measures each packed row, so the output is allocated once at its exact size and the second - // packs every row straight into its place. Collecting packed chunks and then copying them into - // the output held the encoded image twice, with growth slack, at the save's peak (photocraft#6607). + let rb = layout.row_bytes(); + let cs = layout.count_size(); let threads = if cfg!(target_arch = "wasm32") { 1 } else { std::thread::available_parallelism().map_or(1, |v| v.get()).clamp(1, 32) }; - let chunk_rows = rows.div_ceil(threads).max(1); - let row = |r: usize| &decoded[r * rb..(r + 1) * rb]; - let mut lens = vec![0usize; rows]; - for_each_chunk(threads, lens.chunks_mut(chunk_rows).enumerate().collect(), |(c, lens): (usize, &mut [usize])| { - let mut scratch = Vec::with_capacity(rb + rb.div_ceil(128)); - for (i, len) in lens.iter_mut().enumerate() { - scratch.clear(); - packbits::encode(row(c * chunk_rows + i), &mut scratch); - *len = scratch.len(); - } - })?; - let mut out = vec![0u8; rows * cs + lens.iter().sum::()]; - let (counts, mut packed) = out.split_at_mut(rows * cs); - for (count, &n) in counts.chunks_exact_mut(cs).zip(&lens) { - if cs == 2 { - let n = u16::try_from(n).map_err(|_| PsdError::LimitExceeded("RLE row exceeds 65535 bytes (use PSB)"))?; - count.copy_from_slice(&n.to_be_bytes()); - } else { - let n = u32::try_from(n).map_err(|_| PsdError::LimitExceeded("RLE row too large"))?; - count.copy_from_slice(&n.to_be_bytes()); - } + let step = rows.div_ceil(threads).max(1); + let ranges: Vec<_> = (0..rows).step_by(step).map(|start| start..start.saturating_add(step).min(rows)).collect(); + // Scan packets once. Each worker owns its count table and payload; joining copies each + // payload once, trading temporary encoded bytes for the former second packet scan. + let mut chunks: Vec<_> = map_chunks(threads, ranges.clone(), |range| -> Result> { + let table = range.len().checked_mul(cs).ok_or(PsdError::LimitExceeded("RLE table overflow"))?; + let hint = + if cfg!(target_arch = "wasm32") { table } else { table.saturating_add(range.len().saturating_mul(rb.saturating_add(rb.div_ceil(128)))) }; + encode_rle_chunk(decoded, layout, range, hint) + })? + .into_iter() + .collect::>()?; + if chunks.len() == 1 { + return Ok(chunks.remove(0).into_boxed_slice().into_vec()); } - let mut parts = Vec::with_capacity(threads); - for (c, lens) in lens.chunks(chunk_rows).enumerate() { - let (part, rest) = packed.split_at_mut(lens.iter().sum()); - parts.push((c, lens, part)); - packed = rest; + let total = chunks.iter().try_fold(0usize, |n, c| n.checked_add(c.len()).ok_or(PsdError::LimitExceeded("RLE size overflow")))?; + let table = rows.checked_mul(cs).ok_or(PsdError::LimitExceeded("RLE table overflow"))?; + let mut out = Vec::new(); + out.try_reserve_exact(total).map_err(|_| PsdError::OutOfMemory("not enough memory for encoded RLE"))?; + out.resize(table, 0); + for (range, chunk) in ranges.into_iter().zip(chunks) { + let counts = range.len() * cs; + out[range.start * cs..range.end * cs].copy_from_slice(&chunk[..counts]); + out.extend_from_slice(&chunk[counts..]); } - for_each_chunk(threads, parts, |(c, lens, mut part): (usize, &[usize], &mut [u8])| { - let mut scratch = Vec::with_capacity(rb + rb.div_ceil(128)); - for (i, &n) in lens.iter().enumerate() { - scratch.clear(); - packbits::encode(row(c * chunk_rows + i), &mut scratch); - let (dst, rest) = std::mem::take(&mut part).split_at_mut(n); - dst.copy_from_slice(&scratch); - part = rest; - } - })?; Ok(out) } -/// Runs `f` on every item, each on its own thread unless `threads` is 1 (wasm32). -fn for_each_chunk(threads: usize, items: Vec, f: impl Fn(T) + Sync) -> Result<()> { - if threads == 1 { - items.into_iter().for_each(f); - return Ok(()); - } - std::thread::scope(|scope| { - let handles: Vec<_> = items.into_iter().map(|item| scope.spawn(|| f(item))).collect(); - // Join every worker instead of aborting the scope with a thread panic. - let mut ok = true; - for h in handles { - ok &= h.join().is_ok(); +/// Encode one worker's rows, measuring packets only when the spare capacity cannot bound them. +fn encode_rle_chunk(decoded: &[u8], layout: &PlaneLayout, range: std::ops::Range, hint: usize) -> Result> { + const OOM: PsdError = PsdError::OutOfMemory("not enough memory for encoded RLE"); + let rb = layout.row_bytes(); + let cs = layout.count_size(); + let table = range.len().checked_mul(cs).ok_or(PsdError::LimitExceeded("RLE table overflow"))?; + let mut out = Vec::new(); + out.try_reserve_exact(hint.max(table)).or_else(|_| out.try_reserve_exact(table)).map_err(|_| OOM)?; + out.resize(table, 0); + let row_bound = rb.saturating_add(rb.div_ceil(128)); + for (i, row) in range.enumerate() { + let src = &decoded[row * rb..(row + 1) * rb]; + if out.capacity() - out.len() < row_bound { + let needed = packbits::encoded_len(src); + out.try_reserve(needed).or_else(|_| out.try_reserve_exact(needed)).map_err(|_| OOM)?; } - if ok { Ok(()) } else { Err(PsdError::LimitExceeded("an RLE worker panicked")) } - }) + let start = out.len(); + packbits::encode(src, &mut out); + let n = out.len() - start; + let count = &mut out[i * cs..(i + 1) * cs]; + if cs == 2 { + count.copy_from_slice(&u16::try_from(n).map_err(|_| PsdError::LimitExceeded("RLE row exceeds 65535 bytes (use PSB)"))?.to_be_bytes()); + } else { + count.copy_from_slice(&u32::try_from(n).map_err(|_| PsdError::LimitExceeded("RLE row too large"))?.to_be_bytes()); + } + } + Ok(out) +} + +/// Ordered work on the caller's pool, including nested channel and row encoding. +fn map_chunks(threads: usize, items: Vec, f: impl Fn(T) -> R + Sync) -> Result> { + if threads == 1 { + return Ok(items.into_iter().map(f).collect()); + } + #[cfg(not(target_arch = "wasm32"))] + { + use rayon::prelude::*; + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| items.into_par_iter().map(&f).collect())) + .map_err(|_| PsdError::LimitExceeded("an RLE worker panicked")) + } + #[cfg(target_arch = "wasm32")] + Ok(items.into_iter().map(f).collect()) } /// zlib-compresses `data` at the default level. pub fn zip_compress(data: &[u8]) -> Vec { + #[cfg(not(target_arch = "wasm32"))] + let mut e = zip_encoder(); + #[cfg(target_arch = "wasm32")] let mut e = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); - // Writing into a Vec cannot fail. + // The in-memory output always accepts the complete write. let _ = e.write_all(data); - e.finish().unwrap_or_default() + #[cfg(not(target_arch = "wasm32"))] + { + if e.try_finish().is_err() { + return Vec::new(); + } + let output = std::mem::take(e.get_mut()); + // Return only the finished codec; the slot stays empty during compression. + let _ = ZIP_ENCODER.try_with(|slot| slot.set(Some(e))); + output.into_vec() + } + #[cfg(target_arch = "wasm32")] + { + e.finish().unwrap_or_default() + } +} + +#[cfg(not(target_arch = "wasm32"))] +thread_local! { + /// One finished codec per thread; encoded output moves out before it is retained. + static ZIP_ENCODER: std::cell::Cell>> = const { std::cell::Cell::new(None) }; +} + +/// Reset a thread's finished encoder without reallocating its fixed-size working buffers. +/// Nested calls and thread teardown can use a fresh encoder without accessing a busy slot. +#[cfg(not(target_arch = "wasm32"))] +fn zip_encoder() -> flate2::write::ZlibEncoder { + if let Some(mut encoder) = ZIP_ENCODER.try_with(std::cell::Cell::take).ok().flatten() + && encoder.reset(ZipOutput::default()).is_ok() + { + return encoder; + } + flate2::write::ZlibEncoder::new(ZipOutput::default(), flate2::Compression::default()) +} + +/// Bounded native output blocks avoid reallocating large ZIP buffers as they grow. +#[cfg(not(target_arch = "wasm32"))] +#[derive(Default)] +struct ZipOutput { + blocks: Vec>, + current: Vec, +} + +#[cfg(not(target_arch = "wasm32"))] +impl std::io::Write for ZipOutput { + /// Accept all input bytes into bounded output blocks without reserving raw image size. + fn write(&mut self, mut bytes: &[u8]) -> std::io::Result { + let len = bytes.len(); + while !bytes.is_empty() { + if self.current.len() == self.current.capacity() { + if !self.current.is_empty() { + self.blocks.push(std::mem::take(&mut self.current)); + } + self.current = Vec::with_capacity(if self.blocks.is_empty() { 8192 } else { 256 << 10 }); + } + let part = &mut self.current; + let (head, rest) = bytes.split_at((part.capacity() - part.len()).min(bytes.len())); + part.extend_from_slice(head); + bytes = rest; + } + Ok(len) + } + + /// Complete a no-op flush because every write is already retained in memory. + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +#[cfg(not(target_arch = "wasm32"))] +impl ZipOutput { + /// Join encoded ZIP blocks into one vector sized to the compressed payload. + fn into_vec(mut self) -> Vec { + self.blocks.push(self.current); + if self.blocks.len() == 1 { + return self.blocks.pop().unwrap_or_default().into_boxed_slice().into_vec(); + } + let len = self.blocks.iter().fold(0usize, |n, c| n.saturating_add(c.len())); + let mut out = Vec::with_capacity(len); + for part in self.blocks { + out.extend_from_slice(&part); + } + out + } } /// Decompresses a zlib stream, requiring at least `expected` output bytes and @@ -489,6 +641,20 @@ pub fn predict(buf: &mut [u8], layout: &PlaneLayout) -> Result<()> { if rb == 0 { return Ok(()); } + #[cfg(not(target_arch = "wasm32"))] + if buf.len() >= 1 << 20 && matches!(layout.depth, 1 | 8 | 16 | 32) { + use rayon::prelude::*; + let threads = rayon::current_num_threads().min(32); + let rows = buf.len() / rb; + let step = rows.div_ceil(threads).max(1) * rb; + return buf[..rows * rb].par_chunks_mut(step).try_for_each(|part| predict_rows(part, layout)); + } + predict_rows(buf, layout) +} + +/// Apply per-row prediction, reusing scratch for float byte-plane shuffling. +fn predict_rows(buf: &mut [u8], layout: &PlaneLayout) -> Result<()> { + let rb = layout.row_bytes(); let w = layout.width; let mut scratch = Vec::new(); for row in buf.chunks_exact_mut(rb) { @@ -509,7 +675,6 @@ pub fn predict(buf: &mut [u8], layout: &PlaneLayout) -> Result<()> { } } 32 => { - scratch.clear(); scratch.resize(rb, 0); // Split into four byte planes (MSB first) and delta the whole // plane-major row. Plane k's first delta uses the last byte of diff --git a/crates/psd/src/compression/tests.rs b/crates/psd/src/compression/tests.rs index 9529af65..9f8c6166 100644 --- a/crates/psd/src/compression/tests.rs +++ b/crates/psd/src/compression/tests.rs @@ -137,13 +137,84 @@ fn layout(planes: usize, width: usize, height: usize, depth: u16, version: Versi PlaneLayout { planes, width, height, depth, version } } +/// Verify PackBits round trips empty, literal and repeated-byte inputs. fn pb_roundtrip(src: &[u8]) -> Vec { let enc = packbits::encode_vec(src); + assert_eq!(packbits::encoded_len(src), enc.len()); let dec = packbits::decode(&enc, src.len()).unwrap(); assert_eq!(dec, src, "encoded: {enc:?}"); enc } +/// Compare block scanning with scalar packet boundaries over varied runs and literals. +#[test] +fn packbits_block_scan_preserves_scalar_packets() { + /// Encode PackBits with the previous scalar packet scan as a byte oracle. + fn scalar(src: &[u8]) -> Vec { + let mut out = Vec::new(); + let mut i = 0; + while i < src.len() { + let mut run = 1; + while i + run < src.len() && run < 128 && src[i + run] == src[i] { + run += 1; + } + if run >= 2 { + out.extend_from_slice(&[(1i16 - run as i16) as u8, src[i]]); + i += run; + } else { + let start = i; + while i < src.len() && i - start < 128 { + if i + 2 < src.len() && src[i] == src[i + 1] && src[i] == src[i + 2] { + break; + } + i += 1; + } + out.push((i - start - 1) as u8); + out.extend_from_slice(&src[start..i]); + } + } + out + } + let check = |src: &[u8]| { + let expected = scalar(src); + let mut actual = vec![42]; + packbits::encode(src, &mut actual); + assert_eq!(&actual[1..], expected, "source {src:?}"); + assert_eq!(packbits::encoded_len(src), expected.len()); + assert!(expected.len() <= src.len() + src.len().div_ceil(128)); + }; + for n in 0..=10 { + for mut bits in 0..3usize.pow(n) { + let src: Vec<_> = (0..n) + .map(|_| { + let b = [0, 1, 255][bits % 3]; + bits /= 3; + b + }) + .collect(); + check(&src); + } + } + for start in 0..260 { + for run in [2, 3, 7, 8, 9, 127, 128, 129, 256] { + let mut src: Vec<_> = (0..520).map(|i| (i % 251) as u8).collect(); + src[start..start + run].fill(255); + check(&src); + } + } + for pair in [[0, 1], [0, 255], [128, 127], [255, 255]] { + for n in [127, 128, 129, 255, 256, 257, 12000, 12001] { + let mut src: Vec<_> = pair.into_iter().cycle().take(n).collect(); + check(&src); + for i in [0, 7, 8, 127.min(n - 1), 128.min(n - 1), n - 1] { + src[i] ^= 1; + check(&src); + src[i] ^= 1; + } + } + } +} + #[test] fn zip_absurd_declared_size_is_an_error() { let z = zip_compress(&[7u8; 64]); @@ -452,6 +523,24 @@ fn prediction_32bit_floats_roundtrip() { assert_eq!(decode_planes(Compression::ZipPrediction, &enc, &l).unwrap(), v); } +/// Compare fused float prediction with explicit byte-plane shuffling and delta coding. +#[test] +fn fused_prediction_matches_shuffled_byte_deltas() { + for w in [1, 2, 3, 31, 256, 257] { + let l = layout(2, w, 3, 32, Version::Psb); + let mut actual: Vec<_> = (0..l.decoded_len().unwrap()).map(|i| (i * 31 + i / 7) as u8).collect(); + let expected: Vec<_> = actual + .chunks_exact(w * 4) + .flat_map(|row| { + let shuffled: Vec<_> = (0..4).flat_map(|k| row.iter().skip(k).step_by(4).copied()).collect(); + (0..shuffled.len()).map(move |i| shuffled[i].wrapping_sub(if i == 0 { 0 } else { shuffled[i - 1] })) + }) + .collect(); + predict(&mut actual, &l).unwrap(); + assert_eq!(actual, expected, "width {w}"); + } +} + #[test] fn zip_bomb_is_bounded() { let big = zip_compress(&vec![0u8; 1 << 20]); @@ -463,3 +552,206 @@ fn zip_bomb_is_bounded() { fn zip_garbage_errors() { assert!(zip_decompress(&[1, 2, 3, 4], 4).is_err()); } + +/// Compare bounded ZIP output with the original growing writer over varied inputs. +#[test] +fn zip_output_matches_the_growing_writer() { + let mut seed = 0x12345678u32; + let noise = (0..1_048_577) + .map(|_| { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + seed as u8 + }) + .collect(); + for bytes in [Vec::new(), vec![0; 131072], (0..131072).map(|i| (i * 31 + i / 173) as u8).collect(), noise].iter().cycle().take(12) { + let mut old = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + old.write_all(bytes).unwrap(); + let packed = zip_compress(bytes); + assert_eq!(packed, old.finish().unwrap()); + #[cfg(not(target_arch = "wasm32"))] + assert_eq!(packed.capacity(), packed.len()); + } +} + +/// Workers retain only private codec state; a checked-out slot supports nested compression. +#[test] +#[cfg(not(target_arch = "wasm32"))] +fn zip_encoder_reuse_is_private_and_retains_no_output() { + std::thread::scope(|scope| { + let workers: Vec<_> = (0..2) + .map(|byte| { + scope.spawn(move || { + let mut seed = 0x12345678u32; + for len in [1_048_577usize, 32769, 0, 259, 258, 257, 1, 32767, 32768, 2, 262145] { + let bytes: Vec<_> = (0..len) + .map(|_| { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + if byte == 0 { 0 } else { seed as u8 } + }) + .collect(); + let packed = zip_compress(&bytes); + if byte == 1 && len == 1_048_577 { + assert!(packed.len() > 256 << 10, "the retention oracle must exercise multiple output blocks"); + } + let encoder = ZIP_ENCODER.with(std::cell::Cell::take).unwrap(); + assert_eq!(encoder.total_in(), len as u64); + assert_eq!(encoder.total_out(), packed.len() as u64); + assert_eq!(encoder.get_ref().blocks.capacity(), 0); + assert_eq!(encoder.get_ref().current.capacity(), 0); + assert_eq!(zip_decompress(&zip_compress(b"nested"), 6).unwrap(), b"nested"); + ZIP_ENCODER.with(|slot| slot.set(Some(encoder))); + let mut fresh = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + fresh.write_all(&bytes).unwrap(); + assert_eq!(zip_compress(&bytes), fresh.finish().unwrap()); + } + }) + }) + .collect(); + for worker in workers { + worker.join().unwrap(); + } + }); +} + +/// Compression stays byte-identical during teardown, including an unavailable native cache. +#[test] +#[cfg(not(target_arch = "wasm32"))] +fn zip_compression_survives_thread_local_teardown() { + use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }; + struct OnDrop(Arc<(AtomicBool, AtomicBool)>); + impl Drop for OnDrop { + /// Record whether teardown reached the destroyed cache and preserved fresh-stream bytes. + fn drop(&mut self) { + let unavailable = ZIP_ENCODER.try_with(|_| ()).is_err(); + let mut fresh = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + fresh.write_all(b"teardown").unwrap(); + let same = zip_compress(b"teardown") == fresh.finish().unwrap(); + self.0.0.store(same, Ordering::SeqCst); + self.0.1.store(unavailable, Ordering::SeqCst); + } + } + thread_local! { static AFTER_CODEC: std::cell::Cell> = const { std::cell::Cell::new(None) }; } + let checked = Arc::new((AtomicBool::new(false), AtomicBool::new(false))); + let flag = checked.clone(); + std::thread::spawn(move || { + AFTER_CODEC.with(|slot| slot.set(Some(OnDrop(flag)))); + zip_compress(b"first"); + }) + .join() + .unwrap(); + assert!(checked.0.load(Ordering::SeqCst), "teardown must preserve the fresh encoder's bytes"); + // OS-key TLS may destroy keys in another order; native TLS runs these destructors in reverse. + #[cfg(any(target_vendor = "apple", target_os = "linux", target_env = "msvc"))] + assert!(checked.1.load(Ordering::SeqCst), "teardown must exercise the uncached fallback"); +} + +/// Verify multiple writes retain their bytes and respect ZIP output block bounds. +#[test] +#[cfg(not(target_arch = "wasm32"))] +fn zip_output_blocks_stay_bounded_across_writes() { + let bytes = vec![42; (1 << 20) + 1]; + let mut output = ZipOutput::default(); + output.write_all(&[]).unwrap(); + for part in bytes.chunks(32767) { + output.write_all(part).unwrap(); + } + assert!(output.blocks.iter().chain([&output.current]).all(|b| b.capacity() <= 256 << 10)); + let packed = output.into_vec(); + assert_eq!(packed, bytes); + assert_eq!(packed.capacity(), packed.len()); +} + +/// Verify compressible ZIP output retains capacity proportional to encoded size. +#[test] +fn compressible_zip_does_not_retain_a_raw_sized_allocation() { + let raw = vec![0; 1 << 20]; + let packed = zip_compress(&raw); + assert!(packed.len() < 1 << 14); + assert!(packed.capacity() < 1 << 16); + assert_eq!(zip_decompress(&packed, raw.len()).unwrap(), raw); +} + +/// Verify failed RLE allocation hints fall back to incremental growth with identical bytes. +#[test] +fn rle_failed_hint_and_incremental_growth_preserve_bytes() { + for version in [Version::Psd, Version::Psb] { + for width in [0, 1, 127, 128, 129, 1000] { + let l = layout(1, width, 9, 8, version); + for fill in [false, true] { + let bytes: Vec<_> = (0..l.decoded_len().unwrap()).map(|i| if fill { 0 } else { (i * 13 + i / 17) as u8 }).collect(); + let expected = encode_planes(Compression::Rle, &bytes, &l).unwrap(); + for hint in [usize::MAX, 0, l.rows().unwrap() * l.count_size()] { + assert_eq!(encode_rle_chunk(&bytes, &l, 0..l.rows().unwrap(), hint).unwrap(), expected); + } + } + } + } +} + +/// Verify measured packet lengths match encoding over mixed runs and literals. +#[test] +fn counted_packbits_matches_mixed_runs_and_literals() { + for n in 0..520 { + let bytes: Vec<_> = (0..n).map(|i| if i % 23 < 9 { (i / 23) as u8 } else { (i * 31) as u8 }).collect(); + assert_eq!(packbits::encoded_len(&bytes), packbits::encode_vec(&bytes).len()); + } +} + +/// Compare parallel prediction with independent serial rows at each supported depth. +#[test] +fn prediction_matches_independent_serial_rows() { + for depth in [1, 8, 16, 32] { + let layout = layout(4, if depth == 1 { 8193 } else { 257 }, 1031, depth, Version::Psb); + let rb = layout.row_bytes(); + let bytes: Vec<_> = (0..layout.decoded_len().unwrap() + 3).map(|i| (i * 13 + i / 17) as u8).collect(); + let mut expected = bytes.clone(); + for row in expected.chunks_exact_mut(rb) { + let src = row.to_vec(); + if depth == 16 { + for (i, p) in row.as_chunks_mut::<2>().0.iter_mut().enumerate().skip(1) { + let a = u16::from_be_bytes([src[i * 2], src[i * 2 + 1]]); + let b = u16::from_be_bytes([src[i * 2 - 2], src[i * 2 - 1]]); + *p = a.wrapping_sub(b).to_be_bytes(); + } + } else { + let shuffled = if depth == 32 { (0..4).flat_map(|c| src.iter().skip(c).step_by(4).copied()).collect::>() } else { src }; + for (i, p) in row.iter_mut().enumerate() { + *p = shuffled[i].wrapping_sub(if i == 0 { 0 } else { shuffled[i - 1] }); + } + } + } + #[cfg(not(target_arch = "wasm32"))] + for threads in [1, 2, 4] { + use rayon::prelude::*; + let pool = rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap(); + pool.install(|| { + (0..3).into_par_iter().for_each(|_| { + let mut got = bytes.clone(); + predict(&mut got, &layout).unwrap(); + assert_eq!(got, expected, "depth {depth}, pool {threads}"); + }) + }); + } + let mut got = bytes; + predict(&mut got, &layout).unwrap(); + assert_eq!(got, expected); + } +} + +/// Verify RLE chunks use the caller's pool and propagate worker panics as errors. +#[cfg(not(target_arch = "wasm32"))] +#[test] +fn rle_chunks_use_the_callers_pool_and_report_panics() { + let pool = rayon::ThreadPoolBuilder::new().num_threads(2).build().unwrap(); + pool.install(|| { + map_chunks(8, (0..64).collect(), |_| assert!(matches!(rayon::current_thread_index(), Some(0 | 1)))).unwrap(); + assert!(map_chunks(2, vec![0, 1], |_| panic!("worker")).is_err()); + }); +} diff --git a/crates/psd/src/filter_effects.rs b/crates/psd/src/filter_effects.rs index b3d545d7..abda0c50 100644 --- a/crates/psd/src/filter_effects.rs +++ b/crates/psd/src/filter_effects.rs @@ -189,6 +189,52 @@ impl FilterEffects { let mut out = Vec::new(); out.put_u32(self.version); for it in &self.items { + let section = out.begin_len(true); + write_pascal(&mut out, it.id.as_bytes(), 1); + out.put_u32(it.version); + let at = out.begin_len(true); + write_rect(&mut out, &it.rect); + out.put_u32(it.depth); + out.put_u32(it.max_channels); + for i in 0..it.max_channels as usize + 2 { + match it.slots.get(i).and_then(Option::as_ref) { + Some(p) => { + out.put_u32(1); + write_plane(&mut out, p); + } + None => out.put_u32(0), + } + } + // A 64-bit length never overflows. + let _ = out.end_len(at, true); + match &it.mask { + Some((r, p)) => { + out.put_u8(1); + write_rect(&mut out, r); + write_plane(&mut out, p); + } + None => out.put_u8(0), + } + let _ = out.end_len(section, true); + out.resize(out.len().next_multiple_of(4), 0); + } + if self.items.is_empty() { + // One empty section, so the block still parses. + out.put_u64(0); + } + out + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Serialize filter-effect sections through staging buffers as a byte oracle. + fn staged_bytes(fx: &FilterEffects) -> Vec { + let mut out = Vec::new(); + out.put_u32(fx.version); + for it in &fx.items { let mut body = Vec::new(); write_pascal(&mut body, it.id.as_bytes(), 1); body.put_u32(it.version); @@ -219,17 +265,12 @@ impl FilterEffects { out.put(&body); out.resize(out.len().next_multiple_of(4), 0); } - if self.items.is_empty() { + if fx.items.is_empty() { // One empty section, so the block still parses. out.put_u64(0); } out } -} - -#[cfg(test)] -mod tests { - use super::*; fn item(id: &str, mask: bool) -> FilterEffectsItem { let rect = Rect { top: 1, left: 2, bottom: 5, right: 6 }; @@ -240,6 +281,20 @@ mod tests { FilterEffectsItem { id: id.into(), version: 1, rect, depth: 8, max_channels: 24, slots, mask: mask.then(|| (rect, plane(128))) } } + /// Compare direct filter-effect serialization with staged bytes and section lengths. + #[test] + fn direct_sections_match_staged_bytes() { + for version in [1, 2, 3] { + let mut large = item(&"é".repeat(140), true); + large.rect.top = -7; + large.slots[0].as_mut().unwrap().data = vec![73; 8193]; + for items in [Vec::new(), vec![large, item("odd", false), item("four", true)]] { + let fx = FilterEffects { version, items }; + assert_eq!(fx.to_bytes(), staged_bytes(&fx)); + } + } + } + /// The unpadded body of `it`'s section. fn section(it: &FilterEffectsItem) -> Vec { let b = FilterEffects { version: 3, items: vec![it.clone()] }.to_bytes(); diff --git a/crates/psd/src/patterns.rs b/crates/psd/src/patterns.rs index f7fd80a5..23ca8cf3 100644 --- a/crates/psd/src/patterns.rs +++ b/crates/psd/src/patterns.rs @@ -25,6 +25,8 @@ use crate::compression::{Compression, MAX_DECODED_BYTES, PlaneLayout, decode_planes, encode_planes}; use crate::error::{PsdError, Result}; use crate::header::Version; +use crate::io::WriteExt; +use std::borrow::Cow; /// Largest pattern edge accepted (Photoshop's own limit is far lower). const MAX_EDGE: u32 = 30_000; @@ -417,7 +419,11 @@ fn zeroed_pattern_plane(len: usize) -> Result> { Ok(out) } +/// Append one validated pattern, backpatching lengths in its final output buffer. fn write_pattern(p: &PsdPattern, out: &mut Vec) -> Result<()> { + if !matches!(p.depth, 1 | 8 | 16 | 32) { + return Err(PsdError::Unsupported(format!("pattern depth {}", p.depth))); + } if p.width > MAX_EDGE || p.height > MAX_EDGE || p.width > i16::MAX as u32 || p.height > i16::MAX as u32 { return Err(PsdError::LimitExceeded("pattern size")); } @@ -444,46 +450,50 @@ fn write_pattern(p: &PsdPattern, out: &mut Vec) -> Result<()> { out.extend_from_slice(&pal); } let rect = [0i32, 0, p.height as i32, p.width as i32]; - let mut body = Vec::new(); + out.extend_from_slice(&3u32.to_be_bytes()); + let at = out.begin_len(false); for v in rect { - body.extend_from_slice(&v.to_be_bytes()); + out.extend_from_slice(&v.to_be_bytes()); } const SLOTS: u32 = 24; - body.extend_from_slice(&SLOTS.to_be_bytes()); + out.extend_from_slice(&SLOTS.to_be_bytes()); let layout = PlaneLayout { planes: 1, width: p.width as usize, height: p.height as usize, depth: p.depth, version: Version::Psd }; - let array = |plane: &[u8], body: &mut Vec| -> Result<()> { + let array = |plane: &[u8], out: &mut Vec| -> Result<()> { let (comp, data) = if p.depth == 8 { - (1u8, encode_planes(Compression::Rle, plane, &layout)?) + (1u8, Cow::Owned(encode_planes(Compression::Rle, plane, &layout)?)) } else { - (0u8, encode_planes(Compression::Raw, plane, &layout)?) + let total = layout.decoded_len()?; + if plane.len() != total { + return Err(PsdError::invalid(format!("decoded buffer has {} bytes, layout requires {}", plane.len(), total))); + } + (0u8, Cow::Borrowed(plane)) }; - body.extend_from_slice(&1u32.to_be_bytes()); - body.extend_from_slice(&((23 + data.len()) as u32).to_be_bytes()); - body.extend_from_slice(&u32::from(p.depth).to_be_bytes()); + out.extend_from_slice(&1u32.to_be_bytes()); + let len = u32::try_from(data.len()).ok().and_then(|n| n.checked_add(23)).ok_or(PsdError::LimitExceeded("pattern array length"))?; + out.extend_from_slice(&len.to_be_bytes()); + out.extend_from_slice(&u32::from(p.depth).to_be_bytes()); for v in rect { - body.extend_from_slice(&v.to_be_bytes()); + out.extend_from_slice(&v.to_be_bytes()); } - body.extend_from_slice(&p.depth.to_be_bytes()); - body.push(comp); - body.extend_from_slice(&data); + out.extend_from_slice(&p.depth.to_be_bytes()); + out.push(comp); + out.extend_from_slice(&data); Ok(()) }; let nc = p.channels.len().min(SLOTS as usize); for c in &p.channels[..nc] { - array(c, &mut body)?; + array(c, out)?; } for _ in nc..SLOTS as usize { - body.extend_from_slice(&0u32.to_be_bytes()); + out.extend_from_slice(&0u32.to_be_bytes()); } // User mask slot: unused; sheet mask slot: transparency, where Photoshop keeps it (photocraft#8054). - body.extend_from_slice(&0u32.to_be_bytes()); + out.extend_from_slice(&0u32.to_be_bytes()); match &p.alpha { - Some(a) => array(a, &mut body)?, - None => body.extend_from_slice(&0u32.to_be_bytes()), + Some(a) => array(a, out)?, + None => out.extend_from_slice(&0u32.to_be_bytes()), } - out.extend_from_slice(&3u32.to_be_bytes()); - out.extend_from_slice(&(body.len() as u32).to_be_bytes()); - out.extend_from_slice(&body); + out.end_len(at, false)?; Ok(()) } @@ -578,11 +588,10 @@ fn pattern_has_user_mask(r: &mut Rd) -> Result { pub fn write_pattern_block(patterns: &[PsdPattern]) -> Result> { let mut out = Vec::new(); for p in patterns { - let mut one = Vec::new(); - write_pattern(p, &mut one)?; - out.extend_from_slice(&(one.len() as u32).to_be_bytes()); - let pad = (4 - one.len() % 4) % 4; - out.extend_from_slice(&one); + let at = out.begin_len(false); + write_pattern(p, &mut out)?; + out.end_len(at, false)?; + let pad = (4 - (out.len() - at - 4) % 4) % 4; out.extend(std::iter::repeat_n(0u8, pad)); } Ok(out) diff --git a/crates/psd/src/patterns/tests.rs b/crates/psd/src/patterns/tests.rs index d1e27692..ae591bc4 100644 --- a/crates/psd/src/patterns/tests.rs +++ b/crates/psd/src/patterns/tests.rs @@ -17,6 +17,56 @@ fn pat(depth: u16, alpha: bool) -> PsdPattern { } } +/// Compare direct pattern serialization with staged bytes at supported depths. +#[test] +fn direct_writes_match_staged_bytes() { + for depth in [1, 8, 16, 32] { + for alpha in [false, true] { + let mut p = pat(depth, alpha); + if depth == 1 { + p.channels = vec![vec![0xA8; 3]; 3]; + p.alpha = alpha.then(|| vec![0xF8; 3]); + } + p.id = "é".repeat(140); + let mut indexed = p.clone(); + indexed.mode = 2; + indexed.channels.truncate(1); + indexed.palette = Some(vec![73; 17]); + let patterns = [p, indexed]; + let mut block = Vec::new(); + let mut file = b"8BPT\0\x01\0\0\0\x02".to_vec(); + for p in &patterns { + let mut one = Vec::new(); + staged_pattern(p, &mut one).unwrap(); + block.extend_from_slice(&(one.len() as u32).to_be_bytes()); + block.extend_from_slice(&one); + block.resize(block.len().next_multiple_of(4), 0); + file.extend_from_slice(&one); + } + assert_eq!(write_pattern_block(&patterns).unwrap(), block); + assert_eq!(write_pat_file(&patterns).unwrap(), file); + } + } +} + +/// Verify borrowed raw pattern planes retain layout and length validation. +#[test] +fn raw_pattern_writes_still_validate_layout_and_plane_length() { + for depth in [1, 16, 32, 7] { + let p = pat(depth, false); + assert!(write_pattern_block(&[PsdPattern { channels: vec![vec![0; 2]], ..p.clone() }]).is_err()); + assert!(write_pat_file(&[PsdPattern { channels: Vec::new(), alpha: Some(vec![0; 2]), ..p }]).is_err()); + } +} + +/// Verify matching plane lengths cannot bypass unsupported pattern-depth checks. +#[test] +fn writers_reject_unsupported_depth_with_matching_plane_lengths() { + let p = PsdPattern { depth: 7, channels: vec![vec![0; 15]; 3], alpha: Some(vec![0; 15]), ..pat(8, false) }; + assert!(write_pattern_block(std::slice::from_ref(&p)).is_err()); + assert!(write_pat_file(&[p]).is_err()); +} + /// photocraft#9318 review: a name the reader refuses (more than 65,536 UTF-16 units with its NUL) is /// refused on write, so a saved library can't become a file that loads nothing. #[test] @@ -350,3 +400,74 @@ fn zero_size_pattern_rectangles_are_rejected() { set_i32(&mut block, vma + 8, 0); assert!(parse_pattern_block(&block).is_err()); } + +/// Serialize a pattern with the previous staged writer as a byte oracle. +fn staged_pattern(p: &PsdPattern, out: &mut Vec) -> Result<()> { + if p.width > MAX_EDGE || p.height > MAX_EDGE || p.width > i16::MAX as u32 || p.height > i16::MAX as u32 { + return Err(PsdError::LimitExceeded("pattern size")); + } + out.extend_from_slice(&1u32.to_be_bytes()); + out.extend_from_slice(&p.mode.to_be_bytes()); + out.extend_from_slice(&(p.height as u16).to_be_bytes()); + out.extend_from_slice(&(p.width as u16).to_be_bytes()); + let units: Vec = p.name.encode_utf16().chain(std::iter::once(0)).collect(); + // The reader refuses longer names (see `read_pattern`). + if units.len() > 65_536 { + return Err(PsdError::LimitExceeded("pattern name length")); + } + out.extend_from_slice(&(units.len() as u32).to_be_bytes()); + for u in units { + out.extend_from_slice(&u.to_be_bytes()); + } + let id = p.id.as_bytes(); + let idl = id.len().min(255); + out.push(idl as u8); + out.extend_from_slice(&id[..idl]); + if p.mode == 2 { + let mut pal = p.palette.clone().unwrap_or_default(); + pal.resize(768, 0); + out.extend_from_slice(&pal); + } + let rect = [0i32, 0, p.height as i32, p.width as i32]; + let mut body = Vec::new(); + for v in rect { + body.extend_from_slice(&v.to_be_bytes()); + } + const SLOTS: u32 = 24; + body.extend_from_slice(&SLOTS.to_be_bytes()); + let layout = PlaneLayout { planes: 1, width: p.width as usize, height: p.height as usize, depth: p.depth, version: Version::Psd }; + let array = |plane: &[u8], body: &mut Vec| -> Result<()> { + let (comp, data) = if p.depth == 8 { + (1u8, encode_planes(Compression::Rle, plane, &layout)?) + } else { + (0u8, encode_planes(Compression::Raw, plane, &layout)?) + }; + body.extend_from_slice(&1u32.to_be_bytes()); + body.extend_from_slice(&((23 + data.len()) as u32).to_be_bytes()); + body.extend_from_slice(&u32::from(p.depth).to_be_bytes()); + for v in rect { + body.extend_from_slice(&v.to_be_bytes()); + } + body.extend_from_slice(&p.depth.to_be_bytes()); + body.push(comp); + body.extend_from_slice(&data); + Ok(()) + }; + let nc = p.channels.len().min(SLOTS as usize); + for c in &p.channels[..nc] { + array(c, &mut body)?; + } + for _ in nc..SLOTS as usize { + body.extend_from_slice(&0u32.to_be_bytes()); + } + // User mask slot: unused; sheet mask slot: transparency, where Photoshop keeps it (photocraft#8054). + body.extend_from_slice(&0u32.to_be_bytes()); + match &p.alpha { + Some(a) => array(a, &mut body)?, + None => body.extend_from_slice(&0u32.to_be_bytes()), + } + out.extend_from_slice(&3u32.to_be_bytes()); + out.extend_from_slice(&(body.len() as u32).to_be_bytes()); + out.extend_from_slice(&body); + Ok(()) +}