image-codecs: count the coefficients an 8-bit progressive decode keeps

zune-jpeg keeps every coefficient of an 8-bit progressive frame (SOF2)
beside its output: 2 bytes per sample, per component at its sampling.
That is 9.85 GB for a 48,000 x 34,200 4:4:4 file, more than the decoded
image. JpegInfo::progressive_coefficients reports them from the frame
header, so a caller can admit the decode against its real peak.

Co-Authored-By: Νύξ (Nyx, AI) 🌑 <nyx@despotak.is>

crates/image-codecs/src/codecs/jpeg.rs

@@ -40,18 +40,37 @@ pub(crate) struct JpegMeta {
     pub adobe_transform: Option<u8>,
     /// (width, height, components) from the first SOFn.
     pub frame: Option<(u32, u32, u8)>,
+    /// [`crate::JpegInfo::progressive_coefficients`] of the first SOFn.
+    pub progressive_coefficients: u64,
 }
 
 /// The first frame of `b` as [`crate::JpegInfo`], from its markers alone (see [`crate::jpeg_info`]).
 pub(crate) fn info(b: &[u8]) -> Option<crate::JpegInfo> {
-    let (width, height, components) = scan_metadata(b).frame?;
+    let m = scan_metadata(b);
+    let (width, height, components) = m.frame?;
     let layout = match components {
         1 => ChannelLayout::Gray,
         3 => ChannelLayout::Rgb,
         4 => ChannelLayout::Cmyk,
         _ => return None,
     };
-    Some(crate::JpegInfo { width, height, layout, whole_coefficients: uses_craft_jpeg(b) })
+    let progressive_coefficients = m.progressive_coefficients;
+    Some(crate::JpegInfo { width, height, layout, whole_coefficients: uses_craft_jpeg(b), progressive_coefficients })
+}
+
+/// Bytes of the coefficients `zune-jpeg` keeps through an 8-bit progressive frame (SOF2), from
+/// its header segment `seg`: per component, a 64-value block of `i16`s for each of its H × V
+/// blocks in every MCU (A.2), the MCUs covering the frame at the largest sampling factors. 0
+/// for any other frame.
+fn progressive_coefficients(marker: u8, seg: &[u8]) -> u64 {
+    let (0xC2, &[8, h0, h1, w0, w1, n, ref comps @ ..]) = (marker, seg) else { return 0 };
+    let (h, w) = (u64::from(u16::from_be_bytes([h0, h1])), u64::from(u16::from_be_bytes([w0, w1])));
+    let factors: Vec<(u64, u64)> =
+        comps.as_chunks::<3>().0.iter().take(usize::from(n)).map(|&[_, f, _]| (u64::from(f >> 4), u64::from(f & 15))).collect();
+    let hmax = factors.iter().map(|f| f.0).max().unwrap_or(1).max(1);
+    let vmax = factors.iter().map(|f| f.1).max().unwrap_or(1).max(1);
+    let mcus = w.div_ceil(8 * hmax) * h.div_ceil(8 * vmax);
+    factors.iter().map(|&(fh, fv)| mcus * fh * fv * 64 * 2).sum()
 }
 
 /// Walk the marker segments up to SOS and collect metadata.
@@ -112,6 +131,7 @@ pub(crate) fn scan_metadata(b: &[u8]) -> JpegMeta {
                 let h = u16::from_be_bytes([seg[1], seg[2]]) as u32;
                 let w = u16::from_be_bytes([seg[3], seg[4]]) as u32;
                 m.frame = Some((w, h, seg[5]));
+                m.progressive_coefficients = progressive_coefficients(marker, seg);
             }
             0xEE if seg.len() >= 12 && seg.starts_with(b"Adobe") => {
                 m.adobe_transform = Some(seg[11]);
@@ -590,16 +610,45 @@ mod tests {
     fn info_reads_the_frame_and_the_decoder() {
         let img = Image::from_raw(37, 21, ChannelLayout::Rgb, SampleType::U8, vec![128; 37 * 21 * 3]).unwrap();
         let rgb = crate::encode(&img, Format::Jpeg, &EncodeOptions::default()).unwrap();
-        assert_eq!(info(&rgb), Some(crate::JpegInfo { width: 37, height: 21, layout: ChannelLayout::Rgb, whole_coefficients: false }));
+        let rgb_info = crate::JpegInfo { width: 37, height: 21, layout: ChannelLayout::Rgb, whole_coefficients: false, progressive_coefficients: 0 };
+        assert_eq!(info(&rgb), Some(rgb_info));
         // A 12-bit extended frame (SOF1) declaring 40,000 × 30,000 gray, markers only.
         let mut twelve = vec![0xFF, 0xD8, 0xFF, 0xC1, 0, 11, 12];
         twelve.extend_from_slice(&30_000u16.to_be_bytes());
         twelve.extend_from_slice(&40_000u16.to_be_bytes());
         twelve.extend_from_slice(&[1, 1, 0x11, 0, 0xFF, 0xD9]);
-        assert_eq!(info(&twelve), Some(crate::JpegInfo { width: 40_000, height: 30_000, layout: ChannelLayout::Gray, whole_coefficients: true }));
+        let twelve_info =
+            crate::JpegInfo { width: 40_000, height: 30_000, layout: ChannelLayout::Gray, whole_coefficients: true, progressive_coefficients: 0 };
+        assert_eq!(info(&twelve), Some(twelve_info));
         assert_eq!(info(b"not a jpeg"), None);
     }
 
+    /// An 8-bit progressive frame (SOF2) reports every coefficient the decoder keeps through its
+    /// scans: a block of 64 `i16`s per component block of every MCU.
+    #[test]
+    fn info_counts_a_progressive_frames_coefficients() {
+        // Markers only: SOI, a frame of `marker` declaring 48,000 × 34,200 with `factors`, EOI.
+        let frame = |marker: u8, precision: u8, factors: &[u8]| {
+            let mut b = vec![0xFF, 0xD8, 0xFF, marker, 0, 8 + 3 * factors.len() as u8, precision];
+            b.extend_from_slice(&34_200u16.to_be_bytes());
+            b.extend_from_slice(&48_000u16.to_be_bytes());
+            b.push(factors.len() as u8);
+            for (i, &f) in factors.iter().enumerate() {
+                b.extend_from_slice(&[i as u8 + 1, f, 0]);
+            }
+            b.extend_from_slice(&[0xFF, 0xD9]);
+            b
+        };
+        let coefficients = |b: &[u8]| info(b).map(|i| i.progressive_coefficients);
+        // 4:4:4: 6,000 × 4,275 MCUs of one block per component.
+        assert_eq!(coefficients(&frame(0xC2, 8, &[0x11, 0x11, 0x11])), Some(6_000 * 4_275 * 3 * 128));
+        // 4:2:0: 3,000 × 2,138 MCUs of 4 + 1 + 1 blocks.
+        assert_eq!(coefficients(&frame(0xC2, 8, &[0x22, 0x11, 0x11])), Some(3_000 * 2_138 * 6 * 128));
+        // Baseline and 12-bit progressive frames (the latter decoded by craft-jpeg) hold none.
+        assert_eq!(coefficients(&frame(0xC0, 8, &[0x11, 0x11, 0x11])), Some(0));
+        assert_eq!(coefficients(&frame(0xC2, 12, &[0x11, 0x11, 0x11])), Some(0));
+    }
+
     /// An APP1 Extended XMP fragment (GUID, total size, offset, data).
     fn extended_xmp(offset: u32) -> Vec<u8> {
         let mut seg = EXTENDED_XMP_HEADER.to_vec();

crates/image-codecs/src/lib.rs

@@ -213,15 +213,20 @@ pub fn probe_size(bytes: &[u8]) -> Result<(u32, u32), CodecError> {
 }
 
 /// A JPEG's frame as its header declares it, read without decoding pixels: the stored size, the
-/// layout [`decode`] returns, and whether [`decode`] hands the file to its whole-coefficient
+/// layout [`decode`] returns, whether [`decode`] hands the file to its whole-coefficient
 /// decoder (arithmetic-coded, lossless or 12-bit frames), whose working set is several times the
-/// decoded image.
+/// decoded image, and the coefficients it holds beside its output for an 8-bit progressive one.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct JpegInfo {
     pub width: u32,
     pub height: u32,
     pub layout: ChannelLayout,
     pub whole_coefficients: bool,
+    /// Bytes of DCT coefficients [`decode`] holds at once beside its output for an 8-bit
+    /// progressive frame (SOF2): every block of every component, at its own sampling, 64
+    /// 16-bit values each, since the scans refine them in turn. 0 for the frames it decodes a
+    /// line of MCUs at a time.
+    pub progressive_coefficients: u64,
 }
 
 /// [`JpegInfo`] of `bytes`; `None` without a readable frame of 1, 3 or 4 components.