address review: charge pictures already held at their own size in the decoder budget

crates/h264/src/decoder.rs

@@ -451,17 +451,19 @@ impl Decoder {
             chroma_planes: if chroma == 0 { 0 } else { 2 },
             bytes_per_sample: 1,
         };
-        // Pictures waiting for output that already left the DPB are held too (one decode call may
-        // carry many access units). A buffer shared with a DPB entry is counted once.
-        let queued = self.out_queue.iter().filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame))).count();
-        let dpb = sps.max_dpb_frames().max(self.dpb.entries.len());
-        let pictures = (dpb + 1 + self.max_in_flight.max(queued)) as u64;
-        ensure!(
-            shape.fits(pictures, self.memory_budget),
-            "{}x{} pictures with {pictures} in the decoder exceed the memory budget",
-            sps.width(),
-            sps.height()
-        );
+        // Pictures already held count at their own size, not the incoming SPS's: after a
+        // resolution change, larger pictures of the old size can still be in the DPB or waiting
+        // for output (one decode call may carry many access units). A buffer shared with a DPB
+        // entry is counted once.
+        let pic = shape.bytes().unwrap_or(u64::MAX);
+        let held = |frames: &mut dyn Iterator<Item = &Arc<Frame>>| frames.fold(0u64, |n, f| n.saturating_add(f.bytes()));
+        let queued = held(&mut self.out_queue.iter().filter(|o| !self.dpb.entries.iter().any(|e| Arc::ptr_eq(&e.frame, &o.frame))).map(|o| &o.frame));
+        let in_dpb = held(&mut self.dpb.entries.iter().map(|e| &e.frame));
+        let dpb = sps.max_dpb_frames().max(self.dpb.entries.len()) as u64;
+        // The DPB (the larger of its new slots and what it holds now), the picture being decoded,
+        // and what is in flight or queued.
+        let bytes = pic.saturating_mul(dpb).max(in_dpb).saturating_add(pic).saturating_add(pic.saturating_mul(self.max_in_flight as u64).max(queued));
+        ensure!(bytes <= self.memory_budget, "{}x{} pictures with {bytes} bytes in the decoder exceed the memory budget", sps.width(), sps.height());
         Ok(())
     }
 

crates/h264/src/picture.rs

@@ -160,6 +160,14 @@ impl Frame {
         self.rows.len()
     }
 
+    /// Sample bytes this picture holds (luma and both chroma planes), as the decoder's memory
+    /// budget counts a picture.
+    pub fn bytes(&self) -> u64 {
+        (self.width as u64)
+            .saturating_mul(self.height as u64)
+            .saturating_add((self.cwidth as u64).saturating_mul(self.cheight as u64).saturating_mul(2))
+    }
+
     /// Publish MB row `r` (ignored if already published).
     pub fn publish(&self, r: usize, row: FrameRow) {
         let _ = self.rows[r].set(row);

crates/h264/src/synth_tests.rs

@@ -715,3 +715,11 @@ fn output_queue_counts_against_the_memory_budget() {
     let err = dec.decode(&[head, ps.concat()].concat(), 0).expect_err("output queue is over budget");
     assert!(err.to_string().contains("memory budget"), "{err}");
 }
+
+/// A held picture is charged at its own size (luma and both 4:2:0 chroma planes), the same
+/// count the budget uses for the incoming SPS's pictures.
+#[test]
+fn a_held_frame_is_charged_at_its_own_size() {
+    let f = crate::picture::Frame::new(0, 0, 2, 3);
+    assert_eq!(f.bytes(), (2 * 3 * 256 + 2 * 2 * 3 * 64) as u64);
+}