mirror of
https://github.com/storytold/deckcraft.git
synced 2026-10-29 21:16:57 -06:00
Nikita RakitinandCursor
6b071aa053
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>
83 lines
3.2 KiB
Rust
83 lines
3.2 KiB
Rust
//! Memory budget for decoded pictures, shared by the video decoders.
|
|
//!
|
|
//! A per-axis size cap does not bound memory: a decoder holds the current picture, every
|
|
//! reference picture and the pictures still being decoded on worker threads. Each decoder checks
|
|
//! the whole picture store against [`DECODE_MEMORY_BUDGET`] with checked arithmetic before it
|
|
//! allocates for a new picture size, so a hostile header is an error, not an out-of-memory abort.
|
|
|
|
/// Bytes the decoded pictures of one decoder may take: 4 GiB on 64-bit targets, 1 GiB on 32-bit
|
|
/// ones (wasm32 has a 4 GiB address space in all).
|
|
pub const DECODE_MEMORY_BUDGET: u64 = if cfg!(target_pointer_width = "64") { 4 << 30 } else { 1 << 30 };
|
|
|
|
/// The shape of one decoded picture as a decoder allocates it.
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct PictureShape {
|
|
/// Allocated luma samples per row and rows (padding included).
|
|
pub width: u64,
|
|
pub height: u64,
|
|
/// Chroma subsampling shifts (0 or 1); ignored when `chroma_planes` is 0.
|
|
pub shift_x: u32,
|
|
pub shift_y: u32,
|
|
/// Chroma planes (0 for monochrome, 2 otherwise).
|
|
pub chroma_planes: u64,
|
|
/// Bytes per sample (1 for 8-bit buffers, 2 for 16-bit ones).
|
|
pub bytes_per_sample: u64,
|
|
}
|
|
|
|
impl PictureShape {
|
|
/// Bytes of one picture, `None` if it does not fit in a `u64`.
|
|
pub fn bytes(&self) -> Option<u64> {
|
|
let luma = self.width.checked_mul(self.height)?;
|
|
let cw = self.width.div_ceil(1 << self.shift_x.min(1));
|
|
let ch = self.height.div_ceil(1 << self.shift_y.min(1));
|
|
let chroma = cw.checked_mul(ch)?.checked_mul(self.chroma_planes)?;
|
|
luma.checked_add(chroma)?.checked_mul(self.bytes_per_sample)
|
|
}
|
|
|
|
/// Bytes of `pictures` such pictures, `None` on overflow.
|
|
pub fn store_bytes(&self, pictures: u64) -> Option<u64> {
|
|
self.bytes()?.checked_mul(pictures)
|
|
}
|
|
|
|
/// Whether `pictures` such pictures fit in `budget` bytes.
|
|
pub fn fits(&self, pictures: u64, budget: u64) -> bool {
|
|
self.store_bytes(pictures).is_some_and(|b| b <= budget)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn shape(w: u64, h: u64) -> PictureShape {
|
|
PictureShape { width: w, height: h, shift_x: 1, shift_y: 1, chroma_planes: 2, bytes_per_sample: 2 }
|
|
}
|
|
|
|
#[test]
|
|
fn bytes_count_luma_and_subsampled_chroma() {
|
|
assert_eq!(shape(4, 4).bytes(), Some((16 + 2 * 4) * 2));
|
|
// Odd sizes round chroma up.
|
|
assert_eq!(shape(3, 3).bytes(), Some((9 + 2 * 4) * 2));
|
|
let mono = PictureShape { chroma_planes: 0, bytes_per_sample: 1, ..shape(10, 10) };
|
|
assert_eq!(mono.bytes(), Some(100));
|
|
let full = PictureShape { shift_x: 0, shift_y: 0, ..shape(10, 10) };
|
|
assert_eq!(full.bytes(), Some(600));
|
|
}
|
|
|
|
#[test]
|
|
fn exact_budget_boundary() {
|
|
let s = shape(4, 4); // 48 bytes
|
|
assert!(s.fits(10, 480));
|
|
assert!(!s.fits(10, 479));
|
|
assert!(!s.fits(11, 480));
|
|
}
|
|
|
|
#[test]
|
|
fn overflow_is_never_within_budget() {
|
|
let huge = shape(u64::MAX / 2, 3);
|
|
assert_eq!(huge.bytes(), None);
|
|
assert!(!huge.fits(1, u64::MAX));
|
|
assert!(!shape(1 << 31, 1 << 31).fits(u64::MAX, u64::MAX));
|
|
}
|
|
}
|