psd: widen staged RLE data to PSB in short reads, never reading it back whole

crates/psd/src/compression.rs

@@ -228,11 +228,18 @@ pub fn decode_planes(compression: Compression, data: &[u8], layout: &PlaneLayout
 /// Bytes a PSD RLE stream (two-byte row counts) of `rows` rows grows by when written as PSB
 /// (four-byte row counts); an error when `data` is too short to hold the count table.
 pub(crate) fn rle_widening(data: &[u8], rows: usize) -> Result<u64> {
-    let table = rows.checked_mul(2).ok_or(PsdError::LimitExceeded("row count overflow"))?;
-    if data.len() < table {
-        return Err(PsdError::UnexpectedEof { offset: data.len(), needed: table - data.len() });
+    rle_widening_len(data.len() as u64, rows)
+}
+
+/// [`rle_widening`] from the stream's length alone, so data still in its source (or staged on
+/// the scratch disk) is not read to measure it.
+pub(crate) fn rle_widening_len(len: u64, rows: usize) -> Result<u64> {
+    let table = rows.checked_mul(2).ok_or(PsdError::LimitExceeded("row count overflow"))? as u64;
+    if len < table {
+        let offset = usize::try_from(len).unwrap_or(usize::MAX);
+        return Err(PsdError::UnexpectedEof { offset, needed: usize::try_from(table - len).unwrap_or(usize::MAX) });
     }
-    Ok(table as u64)
+    Ok(table)
 }
 
 /// Encodes planar data (excluding the 2-byte compression field).

crates/psd/src/file.rs

@@ -2,7 +2,7 @@
 
 use std::io::{Seek, Write};
 
-use crate::compression::{Compression, MAX_DECODED_BYTES, decode_planes, rle_widening, validate_planes};
+use crate::compression::{Compression, MAX_DECODED_BYTES, decode_planes, rle_widening_len, validate_planes};
 use crate::error::{PsdError, Result};
 use crate::header::{ColorMode, Header, Version};
 use crate::image_data::ImageData;
@@ -384,7 +384,7 @@ impl PsdFile {
         })?;
         match self.merged_widened_rows(widen)? {
             0 => out.put_encoded(&self.image_data.data, self.image_data.stored.as_ref()),
-            rows => out.put_widened_rle(&self.image_data.encoded()?, rows),
+            rows => out.put_widened_encoded(&self.image_data.data, self.image_data.stored.as_ref(), rows),
         }
     }
 
@@ -414,13 +414,13 @@ impl PsdFile {
             for c in &layer.channels {
                 match layer.widened_rows(c, true)? {
                     0 => {}
-                    rows => n = n.saturating_add(rle_widening(&c.encoded()?, rows)?),
+                    rows => n = n.saturating_add(rle_widening_len(c.encoded_len(), rows)?),
                 }
             }
         }
         match self.merged_widened_rows(true)? {
             0 => Ok(n),
-            rows => Ok(n.saturating_add(rle_widening(&self.image_data.encoded()?, rows)?)),
+            rows => Ok(n.saturating_add(rle_widening_len(self.image_data.encoded_len(), rows)?)),
         }
     }
 
@@ -463,7 +463,7 @@ impl PsdFile {
         }
         let merged = match self.merged_widened_rows(widen)? {
             0 => 0,
-            rows => rle_widening(&self.image_data.encoded()?, rows)?,
+            rows => rle_widening_len(self.image_data.encoded_len(), rows)?,
         };
         Ok(n.saturating_add(2).saturating_add(self.image_data.encoded_len()).saturating_add(merged))
     }
@@ -747,6 +747,48 @@ mod write_tests {
         assert!(*source.1.lock().unwrap() <= 7, "read {} bytes at once", source.1.lock().unwrap());
     }
 
+    /// #5333: RLE data still in its source widens to PSB in short reads too, the count table and
+    /// then the packed rows, one consecutive pass, never the whole merged image or channel at once.
+    #[test]
+    fn source_backed_rle_widens_in_short_reads() {
+        struct Reads(Vec<u8>, std::sync::Mutex<(usize, Vec<(u64, usize)>)>);
+        impl crate::source::ReadAt for Reads {
+            fn len(&self) -> u64 {
+                self.0.len() as u64
+            }
+            fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
+                let mut seen = self.1.lock().unwrap();
+                seen.0 = seen.0.max(buf.len());
+                seen.1.push((offset, buf.len()));
+                self.0.as_slice().read_at(offset, buf)
+            }
+        }
+        let file = testgen::small(Version::Psd, Compression::Rle);
+        assert_eq!(file.image_data.compression, Compression::Rle);
+        let bytes = file.to_bytes().unwrap();
+        let mut promoted = file.clone();
+        promoted.promote_to_psb().unwrap();
+        let want = promoted.to_bytes().unwrap();
+        let source = std::sync::Arc::new(Reads(bytes, std::sync::Mutex::new((0, Vec::new()))));
+        let mut lazy = PsdFile::from_source(crate::Source::from_arc(source.clone())).unwrap();
+        assert!(lazy.image_data.stored.is_some(), "the merged image stays in the source");
+        lazy.promote_to_psb().unwrap();
+        *source.1.lock().unwrap() = (0, Vec::new());
+        let mut w = std::io::Cursor::new(Vec::new());
+        lazy.write_to(&mut w).unwrap();
+        assert!(w.into_inner() == want);
+        let seen = source.1.lock().unwrap();
+        assert!(seen.0 <= 7, "read {} bytes at once", seen.0);
+        let merged = lazy.image_data.stored.as_ref().unwrap();
+        let reads: Vec<_> = seen.1.iter().filter(|(at, _)| (merged.offset..merged.offset + merged.len).contains(at)).collect();
+        let mut next = merged.offset;
+        for &&(at, n) in &reads {
+            assert_eq!(at, next, "the merged image is read in one consecutive pass");
+            next += n as u64;
+        }
+        assert_eq!(next, merged.offset + merged.len);
+    }
+
     /// #6607: pixel data goes to the stream as it is, never through a staging copy, so a save
     /// holds the model and a short buffer, not a second copy of the file.
     #[test]

crates/psd/src/io.rs

@@ -275,16 +275,43 @@ impl<'w, W: Write + Seek + ?Sized> Sink<'w, W> {
         let table_len = usize::try_from(crate::compression::rle_widening(data, rows)?).map_err(|_| PsdError::LimitExceeded("row count overflow"))?;
         let (table, packed) = data.split_at_checked(table_len).ok_or_else(|| PsdError::invalid("RLE row-count table out of range"))?;
         for part in table.chunks(STAGE / 2) {
-            self.small(|o| {
-                for count in part.as_chunks::<2>().0 {
-                    o.extend_from_slice(&[0, 0, count[0], count[1]]);
-                }
-                Ok(())
-            })?;
+            self.widen_counts(part)?;
         }
         self.put(packed)
     }
 
+    /// [`Self::put_widened_rle`] for data held in memory (`data`) or, when `stored`, still in its
+    /// source: the count table and then the packed rows are read in [`COPY_CHUNK`]-sized parts
+    /// (#5333), so a staged merged image or channel is never read back whole to widen it.
+    pub(crate) fn put_widened_encoded(&mut self, data: &[u8], stored: Option<&crate::source::Stored>, rows: usize) -> std::io::Result<()> {
+        let Some(s) = stored else { return self.put_widened_rle(data, rows) };
+        let table_len = crate::compression::rle_widening_len(s.len, rows)?;
+        // Whole counts per part; the parts follow each other, so a staged range is read in one pass.
+        let table_part = (COPY_CHUNK / 2).max(1) * 2;
+        let mut at = 0u64;
+        while at < s.len {
+            let n = if at < table_len { (table_len - at).min(table_part) } else { (s.len - at).min(COPY_CHUNK) };
+            let part = s.read_part(at, n)?;
+            if at < table_len {
+                self.widen_counts(&part)?
+            } else {
+                self.put(&part)?
+            }
+            at += n;
+        }
+        Ok(())
+    }
+
+    /// Two-byte RLE row counts as four-byte ones (`00 00` and the count), staged.
+    fn widen_counts(&mut self, counts: &[u8]) -> std::io::Result<()> {
+        self.small(|o| {
+            for count in counts.as_chunks::<2>().0 {
+                o.extend_from_slice(&[0, 0, count[0], count[1]]);
+            }
+            Ok(())
+        })
+    }
+
     /// Appends encoded bytes held in memory (`data`) or, when `stored`, still in their source:
     /// those are copied through in [`COPY_CHUNK`] pieces, so a channel or merged image kept on
     /// disk (an imported file's, or one an exporter spilled to scratch) is never read whole.

crates/psd/src/layer.rs

@@ -3,7 +3,7 @@
 use std::borrow::Cow;
 
 use crate::blend::BlendMode;
-use crate::compression::{Compression, PlaneLayout, decode_planes, encode_planes, rle_widening};
+use crate::compression::{Compression, PlaneLayout, decode_planes, encode_planes, rle_widening_len};
 use crate::error::{PsdError, Result};
 use crate::header::{MAX_DIMENSION, Version};
 use crate::io::{Reader, Sink, WriteExt, append_with, decode_legacy_name, read_pascal, write_pascal};
@@ -692,7 +692,7 @@ impl LayerRecord {
         for c in &self.channels {
             out.put_i16(c.id);
             let rows = self.widened_rows(c, widen)?;
-            let grown = if rows == 0 { 0 } else { rle_widening(&c.encoded()?, rows)? };
+            let grown = if rows == 0 { 0 } else { rle_widening_len(c.encoded_len(), rows)? };
             out.put_len(c.stored_len().saturating_add(grown), version.is_psb())?;
         }
         out.put(b"8BIM");
@@ -783,7 +783,7 @@ impl LayerInfo {
         for l in &self.layers {
             for c in l.channels.iter().filter(|c| c.compression.is_some()) {
                 let rows = l.widened_rows(c, widen)?;
-                let grown = if rows == 0 { 0 } else { rle_widening(&c.encoded()?, rows)? };
+                let grown = if rows == 0 { 0 } else { rle_widening_len(c.encoded_len(), rows)? };
                 n = n.saturating_add(2).saturating_add(c.encoded_len()).saturating_add(grown);
             }
         }
@@ -860,7 +860,7 @@ impl LayerInfo {
                     })?;
                     match l.widened_rows(c, widen)? {
                         0 => out.put_encoded(&c.data, c.stored.as_ref())?,
-                        rows => out.put_widened_rle(&c.encoded()?, rows)?,
+                        rows => out.put_widened_encoded(&c.data, c.stored.as_ref(), rows)?,
                     }
                 }
             }