diff --git a/crates/vp9enc/src/bool.rs b/crates/vp9enc/src/bool.rs --- a/crates/vp9enc/src/bool.rs +++ b/crates/vp9enc/src/bool.rs @@ -7,6 +7,9 @@ pub struct BoolEncoder { range: u32, bottom: u32, bit_count: i32, + /// Start of the run of 0xFF bytes at the end of `out` (`out.len()` when it ends in another + /// byte). A carry only touches that run and the byte before it. + ff_start: usize, } impl Default for BoolEncoder { @@ -17,18 +20,37 @@ impl Default for BoolEncoder { impl BoolEncoder { pub fn new() -> Self { - BoolEncoder { out: Vec::new(), range: 255, bottom: 0, bit_count: 24 } + BoolEncoder { out: Vec::new(), range: 255, bottom: 0, bit_count: 24, ff_start: 0 } } + fn push_byte(&mut self, b: u8) { + if b != 255 { + self.ff_start = self.out.len() + 1; + } + self.out.push(b); + } + + /// Carry into the bytes written so far: the trailing 0xFF run becomes 0x00 and the byte + /// before it grows by one (a carry out of the first byte is dropped). Each 0xFF byte is + /// cleared once, so a long run costs O(1) per byte instead of a rescan for every carry. fn add_one(&mut self) { - let mut i = self.out.len(); - while i > 0 { - i -= 1; - if self.out[i] == 255 { - self.out[i] = 0; - } else { - self.out[i] += 1; - return; + let end = self.out.len(); + let start = self.ff_start.min(end); + if let Some(run) = self.out.get_mut(start..) { + run.fill(0); + } + self.ff_start = end; + if let Some(i) = start.checked_sub(1) + && let Some(b) = self.out.get_mut(i) + { + *b = b.wrapping_add(1); + if *b == 255 { + // The grown byte joins the 0xFF bytes before it. + let mut s = i; + while s > 0 && self.out.get(s - 1) == Some(&255) { + s -= 1; + } + self.ff_start = s; } } } @@ -51,7 +73,7 @@ impl BoolEncoder { self.bottom <<= 1; self.bit_count -= 1; if self.bit_count == 0 { - self.out.push((self.bottom >> 24) as u8); + self.push_byte((self.bottom >> 24) as u8); self.bottom &= (1 << 24) - 1; self.bit_count = 8; } @@ -80,7 +102,7 @@ impl BoolEncoder { v <<= 8; } for _ in 0..4 { - self.out.push((v >> 24) as u8); + self.push_byte((v >> 24) as u8); v <<= 8; } self.out @@ -165,4 +187,73 @@ mod tests { assert!(dec.read(1)); } } + + /// The encoder as it was before the carry run was tracked: a carry rescans the bytes. + struct Scan { + out: Vec, + range: u32, + bottom: u32, + bit_count: i32, + } + impl Scan { + fn add_one(&mut self) { + let mut i = self.out.len(); + while i > 0 { + i -= 1; + if self.out[i] == 255 { + self.out[i] = 0; + } else { + self.out[i] += 1; + return; + } + } + } + fn write(&mut self, bit: bool, prob: u8) { + let split = 1 + (((self.range - 1) * prob as u32) >> 8); + if bit { + self.bottom = self.bottom.wrapping_add(split); + self.range -= split; + } else { + self.range = split; + } + while self.range < 128 { + self.range <<= 1; + if self.bottom & (1 << 31) != 0 { + self.add_one(); + } + self.bottom <<= 1; + self.bit_count -= 1; + if self.bit_count == 0 { + self.out.push((self.bottom >> 24) as u8); + self.bottom &= (1 << 24) - 1; + self.bit_count = 8; + } + } + } + } + + #[test] + fn carry_run_tracking_matches_the_rescan() { + let mut seed = 99u32; + for case in 0..300 { + let mut fast = BoolEncoder::new(); + let mut slow = Scan { out: vec![], range: 255, bottom: 0, bit_count: 24 }; + // Skewed probabilities make long 0xFF runs and carries likely. + let extreme = case % 3; + for _ in 0..4000 { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + let p = match extreme { + 0 => (seed >> 8) as u8 | 1, + 1 => 1 + (seed >> 8) as u8 % 4, + _ => 255 - (seed >> 8) as u8 % 4, + }; + let bit = (seed >> 24) % 256 >= p as u32 || (seed >> 3).is_multiple_of(11); + fast.write(bit, p); + slow.write(bit, p); + assert_eq!(fast.out, slow.out, "case {case}"); + } + } + } } diff --git a/crates/vp9enc/src/lib.rs b/crates/vp9enc/src/lib.rs --- a/crates/vp9enc/src/lib.rs +++ b/crates/vp9enc/src/lib.rs @@ -33,12 +33,17 @@ use probs::Probs; use tables::{AC_QLOOKUP, DC_QLOOKUP}; use transform::Quant; +/// Largest width or height the 16-bit frame size fields can hold. +const MAX_DIMENSION: u32 = 1 << 16; + pub(crate) const BLOCK_8X8: u8 = 3; pub(crate) const BLOCK_16X16: u8 = 6; /// Encoder settings. #[derive(Clone, Debug, PartialEq)] pub struct EncoderConfig { + /// Frame size in samples; the stream stores size - 1 in 16 bits, so values above 65536 are + /// clamped to 65536 by [`Vp9Encoder::new`]. pub width: u32, pub height: u32, /// 1..=100; 100 is lossless. With a target bitrate: the quality of the first frame. @@ -140,7 +145,10 @@ fn quality_to_qidx(q: u8) -> u8 { } impl Vp9Encoder { - pub fn new(cfg: EncoderConfig) -> Self { + pub fn new(mut cfg: EncoderConfig) -> Self { + // The frame header stores width - 1 and height - 1 in 16 bits. + cfg.width = cfg.width.min(MAX_DIMENSION); + cfg.height = cfg.height.min(MAX_DIMENSION); let q = cfg.quality.clamp(1, 100); let lossless = q >= 100 && cfg.target_kbps.is_none(); let base_qidx = if lossless { 0 } else { quality_to_qidx(q) }; @@ -461,3 +469,16 @@ pub fn alpha_to_yuv420(rgba: &[u8], w: u32, h: u32) -> (Vec, Vec, Vec = rgba.as_chunks::<4>().0.iter().take(w * h).map(|p| p[3]).collect(); (y, vec![128; cw * ch], vec![128; cw * ch]) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sizes_beyond_the_16_bit_header_fields_are_clamped() { + let enc = Vp9Encoder::new(EncoderConfig { width: 70_000, height: u32::MAX, ..Default::default() }); + assert_eq!((enc.cfg.width, enc.cfg.height), (65_536, 65_536)); + let enc = Vp9Encoder::new(EncoderConfig { width: 65_536, height: 1, ..Default::default() }); + assert_eq!((enc.cfg.width, enc.cfg.height), (65_536, 1)); + } +} diff --git a/crates/vp9enc/src/probs.rs b/crates/vp9enc/src/probs.rs --- a/crates/vp9enc/src/probs.rs +++ b/crates/vp9enc/src/probs.rs @@ -212,7 +212,32 @@ fn inv_remap_prob(delta: usize, prob: u8) -> u8 { /// The `decode_term_subexp` value that moves `old` to `new`. fn delta_for(old: u8, new: u8) -> Option { - (0..255).find(|&d| inv_remap_prob(d, old) == new) + match delta_table().get(old as usize * 256 + new as usize) { + Some(&d) if d != NO_DELTA => Some(d as usize), + _ => None, + } +} + +/// Marks "no delta reaches this probability" in [`delta_table`]. +const NO_DELTA: u8 = 255; + +/// `delta_for` for every (old, new) pair, indexed `old * 256 + new`: the smallest delta that +/// moves `old` to `new`. Filled from the largest delta down so the smallest one wins. +fn delta_table() -> &'static [u8] { + use std::sync::OnceLock; + static T: OnceLock> = OnceLock::new(); + T.get_or_init(|| { + let mut t = vec![NO_DELTA; 256 * 256]; + for old in 0..=255u8 { + for d in (0..255usize).rev() { + let new = inv_remap_prob(d, old); + if let Some(e) = t.get_mut(old as usize * 256 + new as usize) { + *e = d as u8; + } + } + } + t + }) } /// Bits of `decode_term_subexp` for `d`. @@ -415,6 +440,16 @@ mod tests { } } + #[test] + fn delta_table_matches_the_linear_search() { + for old in 0..=255u8 { + for new in 0..=255u8 { + let slow = (0..255).find(|&d| inv_remap_prob(d, old) == new); + assert_eq!(delta_for(old, new), slow, "{old} -> {new}"); + } + } + } + #[test] fn inv_map_table_shape() { let t = inv_map_table();