Codecs: decode PNM straight into the budgeted output buffer (#2938) (#3291)

Limits::max_alloc budgets the decoded pixel buffer, but the PNM decoder
collected one u32 per sample first (4x an 8-bit image, plus a u16 copy
for 16-bit) and PFM built a Vec<f32> beside the output bytes. Samples are
now rescaled and written straight into the native-endian output buffer
(reserved fallibly), so the only header-sized allocation is the one
Limits::check already charged.

Co-authored-by: Claude Opus 5.5 <[email protected]>

crates/codecs/src/codecs/pnm.rs

@@ -79,17 +79,43 @@ fn sample_for_maxval(maxval: u32) -> Result<SampleType, CodecError> {
     }
 }
 
-/// Build an image from raw integer values in `[0, maxval]`.
-fn from_values(w: u32, h: u32, layout: ChannelLayout, maxval: u32, values: Vec<u32>) -> Result<Image, CodecError> {
-    let sample = sample_for_maxval(maxval)?;
-    if values.iter().any(|&v| v > maxval) {
-        return Err(err("sample exceeds maxval"));
+/// The decoded raster, built straight in the output representation. It is the only buffer sized
+/// by the header, so `Limits::check` on the output budgets all of it: no wider scratch buffer
+/// (such as one `u32` per sample) sits beside it (#2938).
+struct Samples {
+    maxval: u32,
+    sample: SampleType,
+    bytes: Vec<u8>,
+}
+
+impl Samples {
+    /// Reserve room for `cap` samples (callers bound `cap` by the file length where the header
+    /// alone could overstate it).
+    fn with_capacity(maxval: u32, cap: usize) -> Result<Self, CodecError> {
+        let sample = sample_for_maxval(maxval)?;
+        let len = cap.checked_mul(sample.bytes()).ok_or_else(|| CodecError::LimitExceeded(format!("{cap} samples do not fit in memory")))?;
+        let mut bytes = Vec::new();
+        bytes.try_reserve_exact(len).map_err(|_| CodecError::LimitExceeded(format!("cannot allocate {len} bytes for the image")))?;
+        Ok(Samples { maxval, sample, bytes })
     }
-    match (sample, maxval) {
-        (SampleType::U8, 255) => Image::from_u8(w, h, layout, values.into_iter().map(|v| v as u8).collect()),
-        (SampleType::U16, 65535) => Image::from_u16(w, h, layout, &values.into_iter().map(|v| v as u16).collect::<Vec<_>>()),
-        (SampleType::U8, m) => Image::from_u8(w, h, layout, values.into_iter().map(|v| ((v * 255 + m / 2) / m) as u8).collect()),
-        (_, m) => Image::from_u16(w, h, layout, &values.into_iter().map(|v| ((v as u64 * 65535 + m as u64 / 2) / m as u64) as u16).collect::<Vec<_>>()),
+
+    /// Append one raw value in `[0, maxval]`, rescaled to the full range of the output sample.
+    fn push(&mut self, v: u32) -> Result<(), CodecError> {
+        let m = self.maxval;
+        if v > m {
+            return Err(err("sample exceeds maxval"));
+        }
+        match (self.sample, m) {
+            (SampleType::U8, 255) => self.bytes.push(v as u8),
+            (SampleType::U8, m) => self.bytes.push(((v * 255 + m / 2) / m) as u8),
+            (_, 65535) => self.bytes.extend_from_slice(&(v as u16).to_ne_bytes()),
+            (_, m) => self.bytes.extend_from_slice(&(((v as u64 * 65535 + m as u64 / 2) / m as u64) as u16).to_ne_bytes()),
+        }
+        Ok(())
+    }
+
+    fn into_image(self, w: u32, h: u32, layout: ChannelLayout) -> Result<Image, CodecError> {
+        Image::from_raw(w, h, layout, self.sample, self.bytes)
     }
 }
 
@@ -102,18 +128,24 @@ fn sample_count(w: u32, h: u32, channels: usize) -> Result<usize, CodecError> {
         .ok_or_else(|| CodecError::LimitExceeded(format!("{w}x{h}x{channels} samples do not fit in memory")))
 }
 
-fn read_binary(data: &[u8], n: usize, maxval: u32) -> Result<Vec<u32>, CodecError> {
+fn read_binary(data: &[u8], n: usize, maxval: u32) -> Result<Samples, CodecError> {
+    let mut out = Samples::with_capacity(maxval, n)?;
     if maxval < 256 {
-        if data.len() < n {
-            return Err(err("truncated raster"));
-        }
-        Ok(data[..n].iter().map(|&v| v as u32).collect())
-    } else {
-        let Some(len) = n.checked_mul(2).filter(|&len| data.len() >= len) else {
+        let Some(raw) = data.get(..n) else {
             return Err(err("truncated raster"));
         };
-        Ok(data[..len].as_chunks::<2>().0.iter().map(|c| u16::from_be_bytes([c[0], c[1]]) as u32).collect())
+        for &v in raw {
+            out.push(v as u32)?;
+        }
+    } else {
+        let Some(raw) = n.checked_mul(2).and_then(|len| data.get(..len)) else {
+            return Err(err("truncated raster"));
+        };
+        for c in raw.as_chunks::<2>().0 {
+            out.push(u16::from_be_bytes(*c) as u32)?;
+        }
     }
+    Ok(out)
 }
 
 pub(crate) fn decode(bytes: &[u8], limits: &Limits) -> Result<Image, CodecError> {
@@ -163,19 +195,19 @@ pub(crate) fn decode(bytes: &[u8], limits: &Limits) -> Result<Image, CodecError>
             let layout = if matches!(kind, b'2' | b'5') { ChannelLayout::Gray } else { ChannelLayout::Rgb };
             limits.check(w, h, layout, sample_for_maxval(maxval)?)?;
             let n = sample_count(w, h, layout.channels())?;
-            let values = if matches!(kind, b'5' | b'6') {
+            let samples = if matches!(kind, b'5' | b'6') {
                 r.single_ws()?;
                 read_binary(r.rest(), n, maxval)?
             } else {
                 // Each ASCII sample is a digit and a separator (the last may lack one), so the
                 // file, not the header, bounds the reservation (#1110).
-                let mut v = Vec::with_capacity(n.min(r.rest().len() / 2 + 1));
+                let mut v = Samples::with_capacity(maxval, n.min(r.rest().len() / 2 + 1))?;
                 for _ in 0..n {
-                    v.push(r.uint()?);
+                    v.push(r.uint()?)?;
                 }
                 v
             };
-            from_values(w, h, layout, maxval, values)
+            samples.into_image(w, h, layout)
         }
         b'7' => decode_pam(&mut r, limits),
         b'F' | b'f' => {
@@ -195,17 +227,22 @@ pub(crate) fn decode(bytes: &[u8], limits: &Limits) -> Result<Image, CodecError>
                 return Err(err("truncated raster"));
             }
             let le = scale < 0.0;
-            let mut out = vec![0f32; row * h as usize];
-            for y in 0..h as usize {
+            // Write native-endian bytes straight into the output buffer: a separate `Vec<f32>`
+            // would double the budgeted allocation (#2938).
+            let row_bytes = row * 4;
+            let mut out = vec![0u8; row_bytes * h as usize];
+            for (y, src) in data.chunks_exact(row_bytes).take(h as usize).enumerate() {
                 // PFM rows are stored bottom-to-top.
-                let dst = (h as usize - 1 - y) * row;
-                for i in 0..row {
-                    let o = (y * row + i) * 4;
-                    let b4 = [data[o], data[o + 1], data[o + 2], data[o + 3]];
-                    out[dst + i] = if le { f32::from_le_bytes(b4) } else { f32::from_be_bytes(b4) };
+                let dst_start = (h as usize - 1 - y) * row_bytes;
+                let Some(dst) = out.get_mut(dst_start..dst_start + row_bytes) else {
+                    return Err(err("truncated raster"));
+                };
+                for (d, s) in dst.as_chunks_mut::<4>().0.iter_mut().zip(src.as_chunks::<4>().0) {
+                    let v = if le { f32::from_le_bytes(*s) } else { f32::from_be_bytes(*s) };
+                    *d = v.to_ne_bytes();
                 }
             }
-            Image::from_f32(w, h, layout, &out)
+            Image::from_raw(w, h, layout, SampleType::F32, out)
         }
         _ => Err(err("unknown Netpbm subtype")),
     }
@@ -250,8 +287,7 @@ fn decode_pam(r: &mut Reader<'_>, limits: &Limits) -> Result<Image, CodecError>
     };
     limits.check(w, h, layout, sample_for_maxval(maxval)?)?;
     let n = sample_count(w, h, layout.channels())?;
-    let values = read_binary(r.rest(), n, maxval)?;
-    from_values(w, h, layout, maxval, values)
+    read_binary(r.rest(), n, maxval)?.into_image(w, h, layout)
 }
 
 pub(crate) fn encode(src: &Image, plan: Plan, _opts: &EncodeOptions) -> Result<Vec<u8>, CodecError> {

crates/codecs/tests/pnm_peak_memory.rs

@@ -0,0 +1,36 @@
+//! #2938: `Limits::max_alloc` budgets the decoded pixel buffer, so the PNM decoder must not build a
+//! wider scratch copy (it used to collect one `u32` per sample, four times an 8-bit image) next to it.
+//!
+//! This file holds a single test so it runs in its own process: it reads the process's peak RSS
+//! (`VmHWM`, Linux only) and nothing else may allocate meanwhile.
+
+#[cfg(target_os = "linux")]
+fn peak_rss_bytes() -> Option<u64> {
+    let status = std::fs::read_to_string("/proc/self/status").ok()?;
+    let line = status.lines().find(|l| l.starts_with("VmHWM:"))?;
+    let kib: u64 = line.trim_start_matches("VmHWM:").trim().trim_end_matches("kB").trim().parse().ok()?;
+    Some(kib * 1024)
+}
+
+#[cfg(target_os = "linux")]
+#[test]
+fn pnm_decode_peak_memory_stays_near_the_output_size() {
+    use photocraft_codecs::*;
+
+    const SIDE: u32 = 4096; // 16 MiB of 8-bit gray output
+    let out_bytes = (SIDE * SIDE) as u64;
+    let mut file = format!("P5\n{SIDE} {SIDE}\n255\n").into_bytes();
+    file.extend((0..out_bytes).map(|i| (i % 251) as u8));
+    // Writing 5 resets the peak RSS to the current RSS; skip where the kernel doesn't allow it.
+    if std::fs::write("/proc/self/clear_refs", "5").is_err() {
+        return;
+    }
+    let Some(before) = peak_rss_bytes() else { return };
+    let img = decode(&file).unwrap();
+    let after = peak_rss_bytes().unwrap();
+    assert_eq!(img.data().len() as u64, out_bytes);
+    assert_eq!(img.data()[300], (300 % 251) as u8);
+    // The output buffer plus slack; a `u32`-per-sample scratch copy alone would be 4× it.
+    let grown = after.saturating_sub(before);
+    assert!(grown < 2 * out_bytes, "decoding a {out_bytes} byte image raised peak RSS by {grown} bytes");
+}