Merge pull request #2449 from vangel76/vp9enc-tables-carry-dimensions

vp9enc: probability delta table, linear bool-encoder carry, clamp oversized frames

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<u8>,
+        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}");
+            }
+        }
+    }
 }

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.
@@ -139,7 +144,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) };
@@ -460,3 +468,16 @@ pub fn alpha_to_yuv420(rgba: &[u8], w: u32, h: u32) -> (Vec<u8>, Vec<u8>, Vec<u8
     let y: Vec<u8> = 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));
+    }
+}

crates/vp9enc/src/probs.rs

@@ -216,7 +216,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<usize> {
-    (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<Vec<u8>> = 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`.
@@ -419,6 +444,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();