address review: size AV1 repeats by the reference format and after layer selection, budget HEVC queued frames at their own size

crates/av1/src/decoder.rs

@@ -118,6 +118,8 @@ enum Output {
         grain: FilmGrainParams,
         apply_film_grain: bool,
         pts: i64,
+        /// What copying the picture out costs, in the reference's own format.
+        bytes: u64,
     },
 }
 
@@ -132,7 +134,7 @@ impl Output {
     fn take(self) -> Result<Picture> {
         match self {
             Output::Decoded(s) => s.wait().map(|p| Arc::try_unwrap(p).unwrap_or_else(|p| (*p).clone())),
-            Output::Existing { slot, seq, width, height, grain, apply_film_grain, pts } => {
+            Output::Existing { slot, seq, width, height, grain, apply_film_grain, pts, .. } => {
                 let rf = slot.wait()?;
                 Ok(output_picture(&seq, &rf, width, height, &grain, apply_film_grain, pts))
             }
@@ -216,9 +218,6 @@ pub struct Decoder {
     out: VecDeque<Output>,
     /// pts of the temporal unit being decoded.
     pts: i64,
-    /// Bytes of the repeated-frame (`show_existing_frame`) pictures queued by the current
-    /// [`Decoder::decode_pts`] call; each is copied out in full when it is returned.
-    chunk_output_bytes: u64,
 }
 
 /// Whether `pictures` frames of `width` x `height`, padded as `FrameBuf::new` pads them, fit in
@@ -301,7 +300,6 @@ impl Decoder {
             memory_budget: deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET,
             out: VecDeque::new(),
             pts: 0,
-            chunk_output_bytes: 0,
         }
     }
 
@@ -334,7 +332,6 @@ impl Decoder {
     /// [`Decoder::decode`] with a timestamp that the pictures shown by this chunk carry.
     pub fn decode_pts(&mut self, data: &[u8], pts: i64) -> Result<Vec<Picture>> {
         self.pts = pts;
-        self.chunk_output_bytes = 0;
         let parsed = self.parse(data);
         // Hand out what is finished, then report a parse error.
         let mut pics = Vec::new();
@@ -384,8 +381,8 @@ impl Decoder {
         // Output policy (7.18.1 note): with scalability, show one frame per temporal unit, the
         // highest spatial layer present.
         let scalable = self.seq.as_ref().is_some_and(|s| s.op_idc != 0);
-        if !scalable {
-            self.out.extend(out.into_iter().map(|(o, _, _)| o));
+        let kept: Vec<Output> = if !scalable {
+            out.into_iter().map(|(o, _, _)| o).collect()
         } else {
             let mut last: Option<(usize, u32)> = None;
             let mut sel: Vec<Output> = Vec::new();
@@ -405,8 +402,27 @@ impl Decoder {
                     }
                 }
             }
-            self.out.extend(sel);
+            sel
+        };
+        // Every repeat is copied out whole, so the repeats this call keeps (after the layer
+        // selection) must not queue more than half the decoder budget; the first one over it and
+        // everything after is dropped.
+        let mut r = r;
+        let mut repeat_bytes = 0u64;
+        let mut ok = Vec::with_capacity(kept.len());
+        for o in kept {
+            if let Output::Existing { bytes, .. } = &o {
+                repeat_bytes = repeat_bytes.saturating_add(*bytes);
+                if repeat_bytes > self.memory_budget / 2 {
+                    if r.is_ok() {
+                        r = Err(Error::Invalid("repeated frame outputs exceed the decoder memory budget"));
+                    }
+                    break;
+                }
+            }
+            ok.push(o);
         }
+        self.out.extend(ok);
         r
     }
 
@@ -505,20 +521,15 @@ impl Decoder {
             let idx = fh.frame_to_show_map_idx;
             let slot = self.ref_frames[idx].clone().ok_or(Error::Invalid("show_existing_frame of an empty slot"))?;
             let info = self.refs[idx].clone();
-            // Every repeat is copied out whole, so a compact header must not queue more
-            // pictures than half the decoder budget allows.
+            // Sized in the reference's own format: a new sequence header may differ from it.
             let shape = deckcraft_bitstream::budget::PictureShape {
                 width: u64::from(info.upscaled_width),
                 height: u64::from(info.frame_height),
-                shift_x: u32::from(seq.color.subsampling_x),
-                shift_y: u32::from(seq.color.subsampling_y),
-                chroma_planes: if seq.color.mono_chrome { 0 } else { 2 },
+                shift_x: u32::from(info.subsampling_x),
+                shift_y: u32::from(info.subsampling_y),
+                chroma_planes: if info.mono_chrome { 0 } else { 2 },
                 bytes_per_sample: 2,
             };
-            self.chunk_output_bytes = self.chunk_output_bytes.saturating_add(shape.bytes().unwrap_or(u64::MAX));
-            if self.chunk_output_bytes > self.memory_budget / 2 {
-                return Err(Error::Invalid("repeated frame outputs exceed the decoder memory budget"));
-            }
             let o = Output::Existing {
                 slot: slot.clone(),
                 seq: seq.clone(),
@@ -527,6 +538,7 @@ impl Decoder {
                 grain: fh.film_grain.clone(),
                 apply_film_grain: self.apply_film_grain,
                 pts: self.pts,
+                bytes: shape.bytes().unwrap_or(u64::MAX),
             };
             if fh.frame_type == KEY_FRAME as u8 {
                 // reference frame loading process (7.21) then refresh every slot (7.20)
@@ -645,6 +657,9 @@ impl Decoder {
                 ri.lf_deltas = fh.lf.deltas;
                 ri.seg_features = fh.seg.features;
                 ri.grain = fh.film_grain.clone();
+                ri.subsampling_x = job.seq.color.subsampling_x;
+                ri.subsampling_y = job.seq.color.subsampling_y;
+                ri.mono_chrome = job.seq.color.mono_chrome;
                 self.ref_frames[i] = Some(job.slot.clone());
             }
         }

crates/av1/src/header.rs

@@ -360,6 +360,10 @@ pub struct RefInfo {
     pub lf_deltas: LoopFilterDeltas,
     pub seg_features: SegmentationFeatures,
     pub grain: FilmGrainParams,
+    /// The colour format the frame was decoded in (a later sequence header may change it).
+    pub subsampling_x: u8,
+    pub subsampling_y: u8,
+    pub mono_chrome: bool,
 }
 
 /// A parsed uncompressed frame header with its derived variables.

crates/hevc/src/decoder.rs

@@ -404,12 +404,24 @@ impl Decoder {
             chroma_planes: if chroma == 0 { 0 } else { 2 },
             bytes_per_sample: 2,
         };
-        // Output frames waiting to be returned stay allocated after they leave the DPB.
-        let queued = self.out_queue.iter().filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame))).count();
-        let pictures = u64::from(sps.max_dec_pic_buffering) + 1 + self.max_in_flight as u64 + queued as u64;
+        // Output frames waiting to be returned stay allocated after they leave the DPB, at their
+        // own size: a resolution change in this call can leave pictures of the previous SPS.
+        let queued: u64 = self
+            .out_queue
+            .iter()
+            .filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame)))
+            .map(|o| {
+                (o.frame.width as u64)
+                    .saturating_mul(o.frame.height as u64)
+                    .saturating_add((o.frame.cwidth as u64).saturating_mul(o.frame.cheight as u64).saturating_mul(2))
+                    .saturating_mul(2)
+            })
+            .fold(0, u64::saturating_add);
+        let pictures = u64::from(sps.max_dec_pic_buffering) + 1 + self.max_in_flight as u64;
+        let total = shape.store_bytes(pictures).map(|b| b.saturating_add(queued));
         ensure!(
-            shape.fits(pictures, self.memory_budget),
-            "{}x{} pictures with {pictures} in the decoder exceed the memory budget",
+            total.is_some_and(|b| b <= self.memory_budget),
+            "{}x{} pictures with {pictures} in the decoder and {queued} bytes queued exceed the memory budget",
             sps.width,
             sps.height
         );