Video: check the whole picture store against a memory budget (#398)

Per-axis size caps let AV1 accept 65536x65536 frames (over 8 GiB of luma)
and did not count references or pictures in flight. A shared budget module
in deckcraft-bitstream sizes a picture with checked arithmetic, and the
AV1, VP9, HEVC and H.264 decoders refuse a picture size whose current
picture, reference slots and in-flight pictures would exceed 4 GiB (1 GiB
on 32-bit), before allocating. set_memory_budget lowers it for hosts and tests.

Co-authored-by: Cursor <cursoragent@cursor.com>

crates/h264/src/decoder.rs

@@ -95,6 +95,8 @@ pub struct Decoder {
     pool: Option<rayon::ThreadPool>,
     in_flight: VecDeque<FrameRef>,
     max_in_flight: usize,
+    /// Bytes the picture store may take ([`deckcraft_bitstream::budget`]).
+    memory_budget: u64,
     threads: usize,
     /// Draft mode: skip deblocking of non-reference pictures ([`Decoder::set_draft`]).
     draft: bool,
@@ -178,6 +180,7 @@ impl Decoder {
             pool,
             in_flight: VecDeque::new(),
             max_in_flight: threads.max(1) + 2,
+            memory_budget: deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET,
             threads: workers,
             draft: false,
         }
@@ -421,7 +424,34 @@ impl Decoder {
         Ok(())
     }
 
+    /// Cap the bytes of decoded pictures (current, references, in flight) this decoder may hold.
+    pub fn set_memory_budget(&mut self, bytes: u64) {
+        self.memory_budget = bytes;
+    }
+
+    /// Refuse an SPS whose pictures would not fit the memory budget, before anything is allocated.
+    fn check_budget(&self, sps: &Sps) -> Result<()> {
+        let chroma = sps.chroma_array_type();
+        let shape = deckcraft_bitstream::budget::PictureShape {
+            width: u64::from(sps.width()),
+            height: u64::from(sps.height()),
+            shift_x: u32::from(chroma != 3),
+            shift_y: u32::from(chroma == 1),
+            chroma_planes: if chroma == 0 { 0 } else { 2 },
+            bytes_per_sample: 1,
+        };
+        let pictures = (sps.max_dpb_frames() + 1 + self.max_in_flight) as u64;
+        ensure!(
+            shape.fits(pictures, self.memory_budget),
+            "{}x{} pictures with {pictures} in the decoder exceed the memory budget",
+            sps.width(),
+            sps.height()
+        );
+        Ok(())
+    }
+
     fn start_picture(&mut self, sh: &SliceHeader, sps: &Arc<Sps>, pts: i64) -> Result<()> {
+        self.check_budget(sps)?;
         // Activate SPS; a resolution change flushes the DPB.
         let changed = match &self.active_sps {
             None => true,

crates/h264/src/synth_tests.rs

@@ -539,3 +539,43 @@ fn poc_type1_reordering() {
         }
     }
 }
+
+/// One IDR picture of PCM macroblocks.
+fn idr_stream(max_refs: u32) -> Vec<u8> {
+    let mut seed = 5u32;
+    let mbs: Vec<Option<Vec<u8>>> = (0..NMB).map(|_| Some(random_mb(&mut seed))).collect();
+    let mut stream = sps(66, max_refs, false, PocType::T2);
+    stream.extend(pps(false, false));
+    stream.extend(slice_nal(&Pic { idr: true, frame_num: 0, ref_idc: 3, mbs, ..Default::default() }, false));
+    stream
+}
+
+fn decodes_within(stream: &[u8], budget: u64) -> bool {
+    let mut dec = Decoder::with_threads(1);
+    dec.set_memory_budget(budget);
+    let Ok(mut out) = dec.decode(stream, 0) else { return false };
+    out.extend(dec.flush());
+    dec.take_error().is_none() && out.len() == 1
+}
+
+#[test]
+fn picture_store_is_checked_against_the_memory_budget() {
+    let stream = idr_stream(4);
+    let sps_nal = sps(66, 4, false, PocType::T2);
+    let parsed = crate::params::Sps::parse(&deckcraft_bitstream::unescape_rbsp(&sps_nal[5..])).unwrap();
+    // 4:2:0, 8-bit: luma plus two quarter-size chroma planes per picture; the DPB, the current
+    // picture and threads + 2 in flight.
+    let one = (NMB * 256 + 2 * NMB * 64) as u64;
+    let pictures = (parsed.max_dpb_frames() + 1 + 3) as u64;
+    assert!(decodes_within(&stream, pictures * one));
+    assert!(!decodes_within(&stream, pictures * one - 1));
+    assert!(decodes_within(&stream, deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET));
+}
+
+#[test]
+fn largest_accepted_h264_size_is_over_the_default_budget() {
+    // 1024 x 1024 macroblocks (16384 x 16384) with 16 references and 3 in flight.
+    let shape =
+        deckcraft_bitstream::budget::PictureShape { width: 16384, height: 16384, shift_x: 1, shift_y: 1, chroma_planes: 2, bytes_per_sample: 1 };
+    assert!(!shape.fits(16 + 1 + 3, deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET));
+}