Skip to main content

wmv_decoder/
vlc_tree.rs

1//! Bit-by-bit Huffman/VLC decoder (upstream get_vlc2 equivalent strategy).
2//!
3//! We intentionally use a tree (bit traversal) instead of a flat lookup table,
4//! because MSMPEG4/WMV2 DC tables contain code lengths up to 24 bits.
5
6use crate::bitreader::BitReader;
7
8#[derive(Clone, Copy, Default)]
9struct Node {
10    left: Option<usize>,
11    right: Option<usize>,
12    sym: Option<i32>,
13}
14
15/// MSB-first VLC tree.
16#[derive(Clone)]
17pub struct VlcTree {
18    nodes: Vec<Node>,
19}
20
21impl VlcTree {
22    pub fn new() -> Self {
23        VlcTree {
24            nodes: vec![Node::default()],
25        }
26    }
27
28    pub fn insert(&mut self, code: u32, len: u8, sym: i32) {
29        let mut cur = 0usize;
30        for bitpos in (0..len).rev() {
31            let bit_is_one = ((code >> bitpos) & 1) != 0;
32
33            // Avoid holding a mutable reference into `self.nodes` across a `push()`.
34            // (A push may reallocate and would invalidate such a reference.)
35            let next = if !bit_is_one {
36                self.nodes[cur].left
37            } else {
38                self.nodes[cur].right
39            };
40            cur = match next {
41                Some(i) => i,
42                None => {
43                    let i = self.nodes.len();
44                    self.nodes.push(Node::default());
45                    if !bit_is_one {
46                        self.nodes[cur].left = Some(i);
47                    } else {
48                        self.nodes[cur].right = Some(i);
49                    }
50                    i
51                }
52            };
53        }
54        self.nodes[cur].sym = Some(sym);
55    }
56
57    /// Decode a symbol. Returns None on EOF or invalid code path.
58    pub fn decode(&self, br: &mut BitReader<'_>) -> Option<i32> {
59        let mut cur = 0usize;
60        loop {
61            if let Some(sym) = self.nodes[cur].sym {
62                return Some(sym);
63            }
64            let bit = br.read_bit()?;
65            cur = if !bit {
66                self.nodes[cur].left?
67            } else {
68                self.nodes[cur].right?
69            };
70        }
71    }
72}