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 Default for VlcTree {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl VlcTree {
28    pub fn new() -> Self {
29        VlcTree {
30            nodes: vec![Node::default()],
31        }
32    }
33
34    pub fn insert(&mut self, code: u32, len: u8, sym: i32) {
35        let mut cur = 0usize;
36        for bitpos in (0..len).rev() {
37            let bit_is_one = ((code >> bitpos) & 1) != 0;
38
39            // Avoid holding a mutable reference into `self.nodes` across a `push()`.
40            // (A push may reallocate and would invalidate such a reference.)
41            let next = if !bit_is_one {
42                self.nodes[cur].left
43            } else {
44                self.nodes[cur].right
45            };
46            cur = match next {
47                Some(i) => i,
48                None => {
49                    let i = self.nodes.len();
50                    self.nodes.push(Node::default());
51                    if !bit_is_one {
52                        self.nodes[cur].left = Some(i);
53                    } else {
54                        self.nodes[cur].right = Some(i);
55                    }
56                    i
57                }
58            };
59        }
60        self.nodes[cur].sym = Some(sym);
61    }
62
63    /// Decode a symbol. Returns None on EOF or invalid code path.
64    pub fn decode(&self, br: &mut BitReader<'_>) -> Option<i32> {
65        let mut cur = 0usize;
66        loop {
67            if let Some(sym) = self.nodes[cur].sym {
68                return Some(sym);
69            }
70            let bit = br.read_bit()?;
71            cur = if !bit {
72                self.nodes[cur].left?
73            } else {
74                self.nodes[cur].right?
75            };
76        }
77    }
78}