craft-video: DeckCraft decoder fixes (HEVC queued-output budget, VP9 failed superframes, AV1 tile metadata cap, H.264 POC overflow) (#1010)

* craft-video: budget HEVC queued output, drop a failed VP9 superframe's repeats, serial AV1 tiles past a metadata cap

Ported from storytold/deckcraft#1778 and #1841, which landed in DeckCraft's
decoder copies after craft-video took them.

* craft-video: H.264 picture order count no longer overflows on extreme stream values

Wrapping arithmetic through the POC type 0/1/2 paths, frame_num_offset and
the MMCO5 prev_ref_lsb update. Ported from storytold/deckcraft#1913.

---------

Co-authored-by: jayzed80 <jzinkant@gmail.com>

crates/video/src/av1/decode/decoder.rs

@@ -908,6 +908,14 @@ fn decode_frame(job: &FrameJob, pool: &Pool, stats: &mut DecodeStats) -> Result<
     Ok((rf, pic))
 }
 
+/// Bytes all parallel tile states may hold of whole-frame CDEF and restoration metadata; past it
+/// the tiles decode one after another on the frame's own metadata.
+const TILE_METADATA_BUDGET: usize = 256 << 20;
+
+fn tile_metadata_fits(tiles: usize, per_tile: usize) -> bool {
+    tiles.checked_mul(per_tile).is_some_and(|n| n <= TILE_METADATA_BUDGET)
+}
+
 /// Decode the tiles of a frame (tile number, data), in parallel when several; returns the final
 /// CDFs of tile context_update_tile_id.
 fn decode_tiles(pool: &Pool, fs: &mut FrameState, cdfs: &Cdfs, tiles: &[(usize, Vec<u8>)]) -> Result<Option<Box<Cdfs>>> {
@@ -921,7 +929,7 @@ fn decode_tiles(pool: &Pool, fs: &mut FrameState, cdfs: &Cdfs, tiles: &[(usize,
     };
     // Intra block copy reads the current frame through the tile's own buffer: keep whole-frame
     // buffers (tiles one after another) for it.
-    if tiles.len() > 1 && pool.is_parallel() && !fs.fh.allow_intrabc {
+    if tiles.len() > 1 && pool.is_parallel() && !fs.fh.allow_intrabc && tile_metadata_fits(tiles.len(), fs.tile_metadata_bytes()) {
         let mut states: Vec<TileState> = tiles
             .iter()
             .map(|(num, _)| {
@@ -1336,6 +1344,13 @@ mod shown_limit_tests {
 #[cfg(test)]
 mod budget_tests {
     use super::*;
+
+    #[test]
+    fn tile_metadata_copies_are_budgeted() {
+        assert!(tile_metadata_fits(64, 1 << 20));
+        assert!(!tile_metadata_fits(4096, 1 << 20));
+        assert!(!tile_metadata_fits(usize::MAX, 2));
+    }
     use crate::budget::DECODE_MEMORY_BUDGET;
 
     fn color(mono: bool) -> ColorConfig {

crates/video/src/av1/decode/state.rs

@@ -253,6 +253,13 @@ impl FrameState {
         self.lr = t.lr;
     }
 
+    /// Bytes of CDEF and loop-restoration metadata each parallel tile state copies whole.
+    pub fn tile_metadata_bytes(&self) -> usize {
+        let lr = &self.lr;
+        // 1 (type) + 6 (Wiener) + 1 (SGR set) + 2 (SGR weights) bytes per unit.
+        self.cdef.len() + lr.lr_type.iter().map(|v| v.len() * 10).sum::<usize>()
+    }
+
     /// Private buffers for the tile covering mode info rows [r0, r1) x columns [c0, c1); tiles
     /// on the right / bottom frame edge also cover the allocation margins.
     pub fn tile_state(&self, r0: usize, r1: usize, c0: usize, c1: usize) -> TileState {

crates/video/src/h264/decode/slice.rs

@@ -462,21 +462,21 @@ impl PocState {
                 let (prev_msb, prev_lsb) = if sh.idr { (0, 0) } else { (self.prev_ref_msb, self.prev_ref_lsb) };
                 let max_lsb = 1i32 << sps.log2_max_poc_lsb;
                 let lsb = sh.pic_order_cnt_lsb as i32;
-                let msb = if lsb < prev_lsb && prev_lsb - lsb >= max_lsb / 2 {
-                    prev_msb + max_lsb
-                } else if lsb > prev_lsb && lsb - prev_lsb > max_lsb / 2 {
-                    prev_msb - max_lsb
+                let msb = if lsb < prev_lsb && prev_lsb.wrapping_sub(lsb) >= max_lsb / 2 {
+                    prev_msb.wrapping_add(max_lsb)
+                } else if lsb > prev_lsb && lsb.wrapping_sub(prev_lsb) > max_lsb / 2 {
+                    prev_msb.wrapping_sub(max_lsb)
                 } else {
                     prev_msb
                 };
-                let top = msb + lsb;
-                let bottom = top + sh.delta_pic_order_cnt_bottom;
+                let top = msb.wrapping_add(lsb);
+                let bottom = top.wrapping_add(sh.delta_pic_order_cnt_bottom);
                 Poc { top, bottom, msb, frame_num_offset: 0 }
             }
             1 => {
                 let frame_num_offset = self.frame_num_offset(sh.idr, sh.frame_num, sps);
                 let n = sps.offset_for_ref_frame.len() as i32;
-                let mut abs_frame_num = if n != 0 { frame_num_offset + sh.frame_num as i32 } else { 0 };
+                let mut abs_frame_num = if n != 0 { frame_num_offset.wrapping_add(sh.frame_num as i32) } else { 0 };
                 if sh.nal_ref_idc == 0 && abs_frame_num > 0 {
                     abs_frame_num -= 1;
                 }
@@ -484,17 +484,17 @@ impl PocState {
                 if abs_frame_num > 0 {
                     let cycle_cnt = (abs_frame_num - 1) / n;
                     let in_cycle = (abs_frame_num - 1) % n;
-                    let delta_cycle: i32 = sps.offset_for_ref_frame.iter().sum();
+                    let delta_cycle = sps.offset_for_ref_frame.iter().fold(0i32, |a, &o| a.wrapping_add(o));
                     expected = cycle_cnt.wrapping_mul(delta_cycle);
                     for &o in &sps.offset_for_ref_frame[..=in_cycle as usize] {
-                        expected += o;
+                        expected = expected.wrapping_add(o);
                     }
                 }
                 if sh.nal_ref_idc == 0 {
-                    expected += sps.offset_for_non_ref_pic;
+                    expected = expected.wrapping_add(sps.offset_for_non_ref_pic);
                 }
-                let top = expected + sh.delta_pic_order_cnt[0];
-                let bottom = top + sps.offset_for_top_to_bottom_field + sh.delta_pic_order_cnt[1];
+                let top = expected.wrapping_add(sh.delta_pic_order_cnt[0]);
+                let bottom = top.wrapping_add(sps.offset_for_top_to_bottom_field).wrapping_add(sh.delta_pic_order_cnt[1]);
                 Poc { top, bottom, msb: 0, frame_num_offset }
             }
             _ => {
@@ -502,9 +502,9 @@ impl PocState {
                 let temp = if sh.idr {
                     0
                 } else if sh.nal_ref_idc == 0 {
-                    2 * (frame_num_offset + sh.frame_num as i32) - 1
+                    2i32.wrapping_mul(frame_num_offset.wrapping_add(sh.frame_num as i32)).wrapping_sub(1)
                 } else {
-                    2 * (frame_num_offset + sh.frame_num as i32)
+                    2i32.wrapping_mul(frame_num_offset.wrapping_add(sh.frame_num as i32))
                 };
                 Poc { top: temp, bottom: temp, msb: 0, frame_num_offset }
             }
@@ -515,7 +515,7 @@ impl PocState {
         if idr {
             0
         } else if self.prev_frame_num > frame_num {
-            self.prev_frame_num_offset + sps.max_frame_num() as i32
+            self.prev_frame_num_offset.wrapping_add(sps.max_frame_num() as i32)
         } else {
             self.prev_frame_num_offset
         }
@@ -528,7 +528,7 @@ impl PocState {
         if sh.nal_ref_idc != 0 {
             if mmco5 {
                 self.prev_ref_msb = 0;
-                self.prev_ref_lsb = poc.top - poc.frame();
+                self.prev_ref_lsb = poc.top.wrapping_sub(poc.frame());
             } else {
                 self.prev_ref_msb = poc.msb;
                 self.prev_ref_lsb = sh.pic_order_cnt_lsb as i32;
@@ -550,3 +550,108 @@ impl PocState {
         self.prev_frame_num = frame_num;
     }
 }
+
+#[cfg(test)]
+mod poc_tests {
+    use super::*;
+    use craft_media::bitstream::BitWriter;
+
+    fn sps(poc_type: u32) -> Sps {
+        let mut w = BitWriter::new();
+        w.write_bits(66, 8);
+        w.write_bits(0, 8);
+        w.write_bits(30, 8);
+        w.write_ue(0); // sps id
+        w.write_ue(0); // log2_max_frame_num - 4
+        w.write_ue(poc_type);
+        match poc_type {
+            0 => w.write_ue(0),
+            _ => {
+                w.write_bit(false);
+                w.write_se(i32::MAX); // offset_for_non_ref_pic
+                w.write_se(i32::MAX); // offset_for_top_to_bottom_field
+                w.write_ue(1);
+                w.write_se(i32::MAX);
+            }
+        }
+        w.write_ue(1); // max_num_ref_frames
+        w.write_bit(false);
+        w.write_ue(0);
+        w.write_ue(0);
+        w.write_bit(true);
+        w.write_bit(true);
+        w.write_bit(false);
+        w.write_bit(false);
+        w.rbsp_trailing();
+        Sps::parse(&w.finish()).expect("sps")
+    }
+
+    fn header(frame_num: u32, lsb: u32, delta_bottom: i32, delta: [i32; 2], nal_ref_idc: u8) -> SliceHeader {
+        SliceHeader {
+            nal_ref_idc,
+            nal_unit_type: 1,
+            idr: false,
+            first_mb_in_slice: 0,
+            slice_type_raw: 0,
+            slice_type: SliceType::P,
+            pps_id: 0,
+            colour_plane_id: 0,
+            frame_num,
+            field_pic: false,
+            bottom_field: false,
+            idr_pic_id: 0,
+            pic_order_cnt_lsb: lsb,
+            delta_pic_order_cnt_bottom: delta_bottom,
+            delta_pic_order_cnt: delta,
+            redundant_pic_cnt: 0,
+            direct_spatial_mv_pred: false,
+            num_ref_idx_active: [1, 0],
+            ref_pic_list_mod: [Vec::new(), Vec::new()],
+            pred_weight_table: None,
+            no_output_of_prior_pics: false,
+            long_term_reference: false,
+            adaptive_ref_pic_marking: false,
+            mmcos: Vec::new(),
+            cabac_init_idc: 0,
+            slice_qp_delta: 0,
+            sp_for_switch: false,
+            slice_qs_delta: 0,
+            disable_deblocking_filter_idc: 0,
+            slice_alpha_c0_offset_div2: 0,
+            slice_beta_offset_div2: 0,
+            slice_group_change_cycle: 0,
+            header_bits: 0,
+        }
+    }
+
+    /// Offsets and deltas are se(v) values that a stream may set anywhere in the i32 range.
+    #[test]
+    fn poc_does_not_overflow_on_extreme_deltas() {
+        for poc_type in 0..=1 {
+            let sps = sps(poc_type);
+            for nal_ref_idc in [0, 1] {
+                let mut state = PocState::default();
+                for frame_num in 0..4 {
+                    for delta in [i32::MIN, -1, 1, i32::MAX] {
+                        let sh = header(frame_num, 5, delta, [delta, delta], nal_ref_idc);
+                        let poc = state.compute(&sh, &sps);
+                        let _ = poc.frame();
+                        state.update(&sh, &poc);
+                    }
+                }
+            }
+        }
+    }
+
+    /// An MMCO5 reference picture can leave `prev_ref_lsb` far below zero; the next type 0
+    /// picture's MSB comparison (`lsb - prev_lsb`) must not overflow.
+    #[test]
+    fn poc_type0_tolerates_negative_prev_ref_lsb() {
+        let sps = sps(0);
+        let sh = header(1, 5, 0, [0, 0], 1);
+        let mut state = PocState { prev_ref_msb: i32::MAX, prev_ref_lsb: i32::MIN, ..PocState::default() };
+        let poc = state.compute(&sh, &sps);
+        let _ = poc.frame();
+        state.update(&sh, &poc);
+    }
+}

crates/video/src/hevc/decode/decoder.rs

@@ -443,10 +443,36 @@ impl Decoder {
             chroma_planes: if chroma == 0 { 0 } else { 2 },
             bytes_per_sample: 2,
         };
+        let frame_bytes = |f: &Frame| {
+            (f.width as u64)
+                .saturating_mul(f.height as u64)
+                .saturating_add((f.cwidth as u64).saturating_mul(f.cheight as u64).saturating_mul(2))
+                .saturating_mul(2)
+        };
+        // 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)) && !self.in_flight.iter().any(|f| Arc::ptr_eq(f, &o.frame)))
+            .map(|o| frame_bytes(&o.frame))
+            .fold(0, u64::saturating_add);
+        // Frames the DPB holds now count at their own size too (an SPS change flushes them into the
+        // queue before the new pictures are made): the DPB part is the larger of its new slots and that.
+        let held: u64 = self.dpb.entries.iter().map(|e| frame_bytes(&e.frame)).fold(0, u64::saturating_add);
+        // A queued frame still being decoded is one of the in-flight pictures already reserved, so the
+        // queue and the in-flight reservation overlap: charge the larger, not both.
         let pictures = u64::from(sps.max_dec_pic_buffering) + 1 + self.max_in_flight as u64;
+        let slots = shape.store_bytes(u64::from(sps.max_dec_pic_buffering));
+        let flight = shape.store_bytes(self.max_in_flight as u64);
+        let total = shape
+            .store_bytes(1)
+            .zip(slots)
+            .zip(flight)
+            .map(|((one, slots), flight)| slots.max(held).saturating_add(one).saturating_add(flight).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
         );
@@ -722,6 +748,32 @@ mod tests {
         decoder
     }
 
+    #[test]
+    /// Output pictures queued after they left the DPB count against the memory budget.
+    fn queued_output_counts_against_the_memory_budget() {
+        let nal = crate::hevc::decode::synth_tests::sps(true);
+        let sps = Sps::parse(&unescape_rbsp(&nal[6..])).unwrap();
+        let mut decoder = Decoder::with_threads(1);
+        // 96x64 4:2:0 in 16-bit samples: the DPB, the current picture and threads + 2 in flight.
+        let one = (96 * 64 + 2 * 48 * 32) * 2;
+        decoder.set_memory_budget((u64::from(sps.max_dec_pic_buffering) + 1 + decoder.max_in_flight as u64) * one);
+        assert!(decoder.check_budget(&sps).is_ok());
+        let meta = Arc::new(OutputMeta {
+            pts: 0,
+            key: false,
+            crop: (0, 0, 16, 16),
+            full_range: false,
+            colour_primaries: 2,
+            transfer_characteristics: 2,
+            matrix_coefficients: 2,
+            sar: (1, 1),
+            bit_depth: 8,
+            draft: false,
+        });
+        decoder.out_queue.push_back(Output { frame: Arc::new(Frame::new(0, 0, 16, 16, 4, 8, 8)), poc: 0, meta });
+        assert!(decoder.check_budget(&sps).is_err());
+    }
+
     #[test]
     /// A failed decode call does not let completed pictures pile up in the output queue.
     fn failed_decode_bounds_the_output_queue() {

crates/video/src/hevc/decode/synth_tests.rs

@@ -161,7 +161,7 @@ fn vps() -> Vec<u8> {
     nal(32, &w.finish())
 }
 
-fn sps(pcm_loop_filter_disabled: bool) -> Vec<u8> {
+pub(crate) fn sps(pcm_loop_filter_disabled: bool) -> Vec<u8> {
     let mut w = BitWriter::new();
     w.write_bits(0, 4);
     w.write_bits(0, 3);

crates/video/src/vp9/decode/decoder.rs

@@ -172,11 +172,17 @@ impl Decoder {
         if frames.len() > 1 {
             self.stats.superframes += 1;
         }
+        let queued = self.pending.len();
         for f in frames {
             if f.is_empty() {
                 continue;
             }
-            self.decode_frame(f, pts)?;
+            if let Err(e) = self.decode_frame(f, pts) {
+                // A failed sample outputs nothing: drop what its earlier subframes queued so a
+                // later drain does not expand each repeat into its own picture.
+                self.pending.truncate(queued);
+                return Err(e);
+            }
         }
         Ok(self.drain(false))
     }

crates/video/tests/vp9_decode/conformance.rs

@@ -191,3 +191,21 @@ fn keyframes_and_reset() {
         }
     }
 }
+
+/// A superframe whose later subframe fails outputs nothing: the repeat its first subframe queued
+/// (show_existing_frame of slot 0) is dropped, not handed out by a later flush.
+#[test]
+fn failed_superframe_drops_its_queued_repeats() {
+    let f = common::fixture("fade_intra_heavy");
+    let Some((ivf, _)) = common::ensure(f) else { return };
+    let data = std::fs::read(ivf).unwrap();
+    let chunks = common::ivf_frames(&data);
+    let mut dec = craft_video::vp9::decode::Decoder::with_threads(1);
+    dec.decode(chunks[0], 0).unwrap();
+    dec.flush();
+    // Profile 0 show_existing_frame of slot 0, then a frame with a bad frame marker; superframe
+    // index with one-byte sizes for two frames.
+    let sample = [0x88, 0x00, 0xc1, 1, 1, 0xc1];
+    assert!(dec.decode(&sample, 1).is_err());
+    assert!(dec.flush().is_empty());
+}