Merge pull request #1197 from dylanmtaylor/codex/photocraft-psd-zlib-stream

PSD: expose incremental zlib stream encoding

crates/psd/src/compression.rs

@@ -10,6 +10,43 @@ use crate::error::{PsdError, Result};
 use crate::header::{Version, row_bytes};
 use crate::io::Reader;
 
+/// An incremental zlib encoder for PSD/PSB channel streams.
+///
+/// PhotoCraft can encode very large channels a band at a time without keeping the compressed
+/// stream's full input or output beside the pixels. `take_output` drains bytes produced so far;
+/// call [`finish`](Self::finish) once all channel data has been written.
+pub struct ZlibEncoder {
+    inner: flate2::write::ZlibEncoder<Vec<u8>>,
+}
+
+impl ZlibEncoder {
+    /// Starts a zlib stream using the PSD writer's default compression level.
+    pub fn new() -> ZlibEncoder {
+        ZlibEncoder { inner: flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()) }
+    }
+
+    /// Adds decoded or predicted channel bytes to the stream.
+    pub fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
+        self.inner.write_all(bytes).map_err(|error| PsdError::Invalid(format!("zlib: {error}")))
+    }
+
+    /// Takes the compressed bytes produced so far, leaving the stream ready for more input.
+    pub fn take_output(&mut self) -> Vec<u8> {
+        std::mem::take(self.inner.get_mut())
+    }
+
+    /// Finishes the zlib stream and returns its remaining compressed bytes.
+    pub fn finish(self) -> Result<Vec<u8>> {
+        self.inner.finish().map_err(|error| PsdError::Invalid(format!("zlib: {error}")))
+    }
+}
+
+impl Default for ZlibEncoder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 /// Maximum number of decoded bytes a single decode call may produce: 8 GiB on
 /// 64-bit targets, 2 GiB elsewhere.
 ///

crates/psd/src/compression/tests.rs

@@ -1,5 +1,21 @@
 use super::*;
 
+#[test]
+fn incremental_zlib_encoder_roundtrips_across_drains() {
+    let expected = (0usize..32_768).map(|n| (n.wrapping_mul(37) % 251) as u8).collect::<Vec<_>>();
+    let mut encoder = ZlibEncoder::new();
+    let mut encoded = Vec::new();
+    for chunk in expected.chunks(997) {
+        encoder.write_all(chunk).unwrap();
+        encoded.extend(encoder.take_output());
+    }
+    encoded.extend(encoder.finish().unwrap());
+
+    let mut decoded = Vec::new();
+    flate2::read::ZlibDecoder::new(encoded.as_slice()).read_to_end(&mut decoded).unwrap();
+    assert_eq!(decoded, expected);
+}
+
 /// The straightforward index-based versions, kept as the oracle for the
 /// register-carried rewrites in [`predict`] and [`unpredict`].
 mod reference {