pito_codecs/vp9/enc/
bool.rs

1use std::sync::LazyLock;
2
3pub struct BoolWriter {
4    pub buf: Vec<u8>,
5    low: u32,
6    range: u32,
7    count: i32,
8}
9
10impl BoolWriter {
11    pub fn new(buf: Vec<u8>) -> Self {
12        let mut w = BoolWriter {
13            buf,
14            low: 0,
15            range: 255,
16            count: -24,
17        };
18        w.put(false, 128);
19        w
20    }
21
22    #[inline(always)]
23    pub fn put(&mut self, bit: bool, prob: u8) {
24        let split = 1 + (((self.range - 1) * prob as u32) >> 8);
25        let mut low = self.low;
26        let mut range = split;
27        if bit {
28            low += split;
29            range = self.range - split;
30        }
31        let mut shift = range.leading_zeros() as i32 - 24;
32        range <<= shift;
33        self.count += shift;
34        if self.count >= 0 {
35            let offset = shift - self.count;
36            if (low << (offset - 1)) & 0x8000_0000 != 0 {
37                let mut x = self.buf.len();
38                while x > 0 && self.buf[x - 1] == 0xff {
39                    self.buf[x - 1] = 0;
40                    x -= 1;
41                }
42                if x > 0 {
43                    self.buf[x - 1] += 1;
44                }
45            }
46            self.buf.push((low >> (24 - offset)) as u8);
47            low <<= offset;
48            shift = self.count;
49            low &= 0xff_ffff;
50            self.count -= 8;
51        }
52        low <<= shift;
53        self.low = low;
54        self.range = range;
55    }
56
57    pub fn literal(&mut self, v: u32, bits: u32) {
58        for i in (0..bits).rev() {
59            self.put((v >> i) & 1 == 1, 128);
60        }
61    }
62
63    pub fn tree(&mut self, tree: &[i8], probs: &[u8], value: u8) {
64        let (bits, len) = tree_path(tree, value);
65        let mut node = 0usize;
66        for i in (0..len).rev() {
67            let b = (bits >> i) & 1;
68            self.put(b == 1, probs[node >> 1]);
69            let next = tree[node + b as usize];
70            if next <= 0 {
71                break;
72            }
73            node = next as usize;
74        }
75    }
76
77    pub fn finish(mut self) -> Vec<u8> {
78        for _ in 0..32 {
79            self.put(false, 128);
80        }
81        if self.buf.last().is_some_and(|&b| b & 0xe0 == 0xc0) {
82            self.buf.push(0);
83        }
84        self.buf
85    }
86}
87
88pub fn tree_path(tree: &[i8], value: u8) -> (u32, u32) {
89    fn walk(tree: &[i8], node: usize, value: u8, bits: u32, len: u32) -> Option<(u32, u32)> {
90        for b in 0..2u32 {
91            let t = tree[node + b as usize];
92            let (nb, nl) = ((bits << 1) | b, len + 1);
93            if t <= 0 {
94                if (-t) as u8 == value {
95                    return Some((nb, nl));
96                }
97            } else if let Some(r) = walk(tree, t as usize, value, nb, nl) {
98                return Some(r);
99            }
100        }
101        None
102    }
103    walk(tree, 0, value, 0, 0).unwrap_or((0, 0))
104}
105
106pub static COST: LazyLock<[u16; 257]> = LazyLock::new(|| {
107    let mut t = [0u16; 257];
108    t[0] = 255 * 16;
109    for (i, c) in t.iter_mut().enumerate().skip(1) {
110        *c = (-(i as f64 / 256.0).log2() * 256.0).round().min(4095.0) as u16;
111    }
112    t
113});
114
115#[inline(always)]
116pub fn cost(bit: bool, prob: u8) -> u32 {
117    COST[if bit {
118        256 - prob as usize
119    } else {
120        prob as usize
121    }] as u32
122}
123
124pub fn tree_costs(tree: &[i8], probs: &[u8], out: &mut [u32]) {
125    fn walk(tree: &[i8], probs: &[u8], node: usize, acc: u32, out: &mut [u32]) {
126        for b in 0..2usize {
127            let t = tree[node + b];
128            let c = acc + cost(b == 1, probs[node >> 1]);
129            if t <= 0 {
130                if let Some(o) = out.get_mut((-t) as usize) {
131                    *o = c;
132                }
133            } else {
134                walk(tree, probs, t as usize, c, out);
135            }
136        }
137    }
138    walk(tree, probs, 0, 0, out);
139}
pito_codecs/vp9/enc/bool.rs

pito_codecs/vp9/enc/
bool.rs
