PSD save: stream Lr16/Lr32 layer blocks into the output (#1158)

Saving a 16- or 32-bit PSD built the Lr16/Lr32 layer payload in a temporary buffer and cloned
every global tagged block before copying them into the file. The payload is now written
straight into the output with its length patched afterwards, and the neighbouring blocks are
borrowed. Output is byte-identical: all 735 corpus PSD/PSBs re-save to the same bytes. A 24 MP
layer saves in 66 ms instead of 75 ms at 16-bit and 132 ms instead of 164 ms at 32-bit, and the
temporary 144-288 MB copy is gone.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

crates/psd/examples/bench_layer_block.rs

@@ -0,0 +1,39 @@
+//! Serialization of one 24 MP raw RGB layer in Lr16/Lr32, without compression or disk I/O.
+//! Run with `cargo run --release -p photocraft-psd --example bench_layer_block`.
+use photocraft_psd::*;
+use std::hint::black_box;
+use std::time::Instant;
+
+fn main() -> Result<()> {
+    for depth in [16, 32] {
+        let mut file = testgen::small(Version::Psd, Compression::Raw);
+        file.header = Header::new(Version::Psd, 6000, 4000, 3, depth, ColorMode::Rgb);
+        let plane_bytes = 6000 * 4000 * usize::from(depth / 8);
+        let mut layer = LayerRecord { rect: Rect::from_xywh(0, 0, 6000, 4000), name: b"RGB".to_vec(), ..Default::default() };
+        for id in 0..3 {
+            layer.channels.push(ChannelData { id, compression: Some(Compression::Raw), data: vec![0; plane_bytes] });
+        }
+        file.layer_info = Some(LayerInfo { merged_alpha: false, layers: vec![layer], padding: None });
+        file.layer_info_placement =
+            LayerInfoPlacement::GlobalBlock { index: 0, signature: *b"8BIM", key: if depth == 16 { *b"Lr16" } else { *b"Lr32" }, padding: None };
+        file.image_data = ImageData { compression: Compression::Raw, data: vec![0; plane_bytes * 3] };
+        black_box(file.to_bytes()?);
+        let mut samples = Vec::new();
+        let mut size = 0;
+        for _ in 0..11 {
+            let t = Instant::now();
+            let bytes = file.to_bytes()?;
+            samples.push(t.elapsed().as_secs_f64() * 1000.0);
+            size = bytes.len();
+            black_box(bytes);
+        }
+        samples.sort_by(f64::total_cmp);
+        println!(
+            "{depth}-bit, {} bytes: p50={:.3} ms p95={:.3} ms",
+            size,
+            samples.get(5).copied().unwrap_or_default(),
+            samples.get(10).copied().unwrap_or_default()
+        );
+    }
+    Ok(())
+}

crates/psd/src/file.rs

@@ -6,7 +6,7 @@ use crate::image_data::ImageData;
 use crate::io::{Reader, WriteExt};
 use crate::layer::{LayerInfo, LayerRecord};
 use crate::resources::{self, ImageResource, ResourceData, ids};
-use crate::tagged::{TaggedBlock, read_blocks, write_blocks};
+use crate::tagged::{TaggedBlock, read_blocks, uses_long_length, write_blocks};
 
 /// Where the layer info was stored in the file.
 #[derive(Debug, Clone, PartialEq, Eq, Default)]
@@ -155,25 +155,34 @@ impl PsdFile {
                 info.write_body(&mut out, v)?;
             }
             out.end_len(li, psb)?;
-            let blocks: std::borrow::Cow<'_, [TaggedBlock]> = match (&self.layer_info, &self.layer_info_placement) {
-                (Some(info), LayerInfoPlacement::GlobalBlock { index, signature, key, padding }) => {
-                    let mut data = Vec::new();
-                    info.write_body(&mut data, v)?;
-                    let mut blocks = self.global_blocks.clone();
-                    let idx = (*index).min(blocks.len());
-                    blocks.insert(idx, TaggedBlock { signature: *signature, key: *key, data, padding: padding.clone() });
-                    blocks.into()
-                }
-                _ => (&self.global_blocks[..]).into(),
-            };
-            let blocks_ref: &[TaggedBlock] = &blocks;
-            let need_mask_field = self.global_layer_mask.is_some() || !blocks_ref.is_empty();
+            let need_mask_field = self.global_layer_mask.is_some() || !self.global_blocks.is_empty() || (in_block && self.layer_info.is_some());
             if need_mask_field {
                 let gm = self.global_layer_mask.as_ref().map_or(&[][..], |g| &g.data[..]);
                 out.put_len(gm.len() as u64, false)?;
                 out.put(gm);
             }
-            write_blocks(&mut out, blocks_ref, v)?;
+            match (&self.layer_info, &self.layer_info_placement) {
+                (Some(info), LayerInfoPlacement::GlobalBlock { index, signature, key, padding }) => {
+                    let (before, after) = self.global_blocks.split_at((*index).min(self.global_blocks.len()));
+                    write_blocks(&mut out, before, v)?;
+                    // Write into the final buffer: a temporary LayerInfo payload and a clone
+                    // of all neighboring opaque blocks otherwise double the save's memory.
+                    out.put(signature);
+                    out.put(key);
+                    let long = uses_long_length(v, key);
+                    let at = out.begin_len(long);
+                    let start = out.len();
+                    info.write_body(&mut out, v)?;
+                    let len = out.len() - start;
+                    out.end_len(at, long)?;
+                    match padding {
+                        Some(bytes) => out.put(bytes),
+                        None => out.extend(std::iter::repeat_n(0u8, len % 2)),
+                    }
+                    write_blocks(&mut out, after, v)?;
+                }
+                _ => write_blocks(&mut out, &self.global_blocks, v)?,
+            }
             out.put(&self.layer_mask_trailing);
         }
         out.end_len(at, psb)?;
@@ -268,3 +277,41 @@ impl PsdFile {
         std::fs::write(path, b)
     }
 }
+
+#[cfg(test)]
+mod write_tests {
+    use super::*;
+    use crate::{Compression, testgen};
+
+    #[test]
+    fn streamed_layer_info_equals_a_materialized_tagged_block() {
+        for version in [Version::Psd, Version::Psb] {
+            for key in [*b"Lr16", *b"Lr32", *b"Layr"] {
+                for index in [0, 1, 2, usize::MAX] {
+                    for padding in [None, Some(vec![]), Some(vec![42]), Some(vec![1, 2, 3])] {
+                        let mut file = testgen::small(version, Compression::Raw);
+                        file.global_blocks = vec![TaggedBlock::new(*b"zz01", vec![11]), TaggedBlock::new(*b"zz02", vec![12, 13])];
+                        file.layer_info_placement = LayerInfoPlacement::GlobalBlock { index, signature: *b"8B64", key, padding: padding.clone() };
+                        // An odd body catches padding calculated against the final file offset.
+                        file.layer_info.as_mut().unwrap().padding = Some(vec![77]);
+                        let actual = file.to_bytes().unwrap();
+                        let mut materialized = file.clone();
+                        let mut data = Vec::new();
+                        materialized.layer_info.take().unwrap().write_body(&mut data, version).unwrap();
+                        materialized.global_blocks.insert(index.min(materialized.global_blocks.len()), TaggedBlock { signature: *b"8B64", key, data, padding });
+                        materialized.layer_info_placement = LayerInfoPlacement::Section;
+                        assert_eq!(actual, materialized.to_bytes().unwrap(), "{version:?} {key:?} {index}");
+                    }
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn global_layer_info_serialization_propagates_layer_limit_errors() {
+        let mut file = testgen::small(Version::Psd, Compression::Raw);
+        file.layer_info = Some(LayerInfo { layers: vec![LayerRecord::default(); 32768], ..Default::default() });
+        file.layer_info_placement = LayerInfoPlacement::GlobalBlock { index: usize::MAX, signature: *b"8BIM", key: *b"Lr16", padding: None };
+        assert!(matches!(file.to_bytes(), Err(PsdError::LimitExceeded(_))));
+    }
+}

crates/psd/tests/layer_block_order.rs

@@ -0,0 +1,26 @@
+//! The global layer block retains its insertion position and neighboring opaque blocks.
+use photocraft_psd::*;
+
+#[test]
+fn global_layer_block_roundtrips_with_opaque_neighbors() {
+    for version in [Version::Psd, Version::Psb] {
+        for depth in [16, 32] {
+            for index in 0..=2 {
+                for signature in [*b"8BIM", *b"8B64"] {
+                    let mut file = testgen::layered(version, ColorMode::Rgb, depth, Compression::Raw);
+                    file.global_blocks = vec![TaggedBlock::new(*b"zz01", vec![7, 8, 9]), TaggedBlock::new(*b"zz02", vec![10, 11])];
+                    file.layer_info_placement = LayerInfoPlacement::GlobalBlock {
+                        index,
+                        signature,
+                        key: if depth == 16 { *b"Lr16" } else { *b"Lr32" },
+                        padding: Some(vec![31, 32, 33]),
+                    };
+                    let bytes = file.to_bytes().unwrap();
+                    let parsed = PsdFile::from_bytes(&bytes).unwrap();
+                    assert_eq!(parsed, file, "{version:?} {depth} {index} {signature:?}");
+                    assert_eq!(parsed.to_bytes().unwrap(), bytes);
+                }
+            }
+        }
+    }
+}