Skip to main content

wmv_decoder/
bitreader.rs

1/// VC-1 / WMV9 Bitstream Reader
2
3#[derive(Clone)]
4pub struct BitReader<'a> {
5    data: &'a [u8],
6    byte_pos: usize,
7    bit_pos: u8,
8    current: u8,
9}
10
11impl<'a> BitReader<'a> {
12    pub fn new(data: &'a [u8]) -> Self {
13        let current = data.first().copied().unwrap_or(0);
14        BitReader {
15            data,
16            byte_pos: 0,
17            bit_pos: 8,
18            current,
19        }
20    }
21
22    /// Create a reader positioned at an arbitrary bit offset (MSB-first).
23    ///
24    /// `bit_offset=0` is identical to `BitReader::new(data)`.
25    pub fn new_at(data: &'a [u8], bit_offset: usize) -> Self {
26        let byte_pos = bit_offset / 8;
27        let bit_in_byte = (bit_offset % 8) as u8;
28        // bit_pos = how many bits remain unread in the current byte.
29        let bit_pos = if bit_in_byte == 0 { 8 } else { 8 - bit_in_byte };
30        let current = data.get(byte_pos).copied().unwrap_or(0);
31        BitReader {
32            data,
33            byte_pos,
34            bit_pos,
35            current,
36        }
37    }
38
39    pub fn read_bits(&mut self, mut n: u8) -> Option<u32> {
40        if n == 0 {
41            return Some(0);
42        }
43        let mut result = 0u32;
44        while n > 0 {
45            if self.byte_pos >= self.data.len() {
46                return None;
47            }
48            let avail = self.bit_pos.min(n);
49            let shift = self.bit_pos - avail;
50            let mask = ((1u16 << avail) - 1) as u8;
51            result = (result << avail) | ((self.current >> shift) & mask) as u32;
52            self.bit_pos -= avail;
53            n -= avail;
54            if self.bit_pos == 0 {
55                self.byte_pos += 1;
56                self.current = self.data.get(self.byte_pos).copied().unwrap_or(0);
57                self.bit_pos = 8;
58            }
59        }
60        Some(result)
61    }
62
63    #[inline]
64    pub fn read_bit(&mut self) -> Option<bool> {
65        self.read_bits(1).map(|b| b != 0)
66    }
67
68    pub fn read_bits_signed(&mut self, n: u8) -> Option<i32> {
69        let v = self.read_bits(n)? as i32;
70        if n == 0 {
71            return Some(0);
72        }
73        let sign = 1i32 << (n - 1);
74        Some(if v & sign != 0 { v - (sign << 1) } else { v })
75    }
76
77    /// Peek up to 24 bits without advancing.
78    pub fn peek_bits(&self, n: u8) -> Option<u32> {
79        if n == 0 {
80            return Some(0);
81        }
82        let mut result = 0u32;
83        let mut bits_left = n;
84        let mut bpos = self.byte_pos;
85        let mut boff = self.bit_pos;
86
87        while bits_left > 0 {
88            let avail = boff.min(bits_left);
89            let shift = boff - avail;
90            let mask = ((1u16 << avail) - 1) as u8;
91            let byte = if bpos < self.data.len() {
92                self.data[bpos]
93            } else {
94                0
95            };
96            result = (result << avail) | ((byte >> shift) & mask) as u32;
97            bits_left -= avail;
98            boff -= avail;
99            if boff == 0 {
100                bpos += 1;
101                boff = 8;
102            }
103        }
104        Some(result)
105    }
106
107    #[inline]
108    pub fn skip_bits(&mut self, n: u8) {
109        let _ = self.read_bits(n);
110    }
111
112    /// Skip an arbitrary number of bits.
113    #[inline]
114    pub fn skip_bits_usize(&mut self, mut n: usize) {
115        while n >= 32 {
116            let _ = self.read_bits(32);
117            n -= 32;
118        }
119        if n > 0 {
120            let _ = self.read_bits(n as u8);
121        }
122    }
123
124    pub fn read_ue(&mut self) -> Option<u32> {
125        let mut zeros = 0u8;
126        while !self.read_bit()? {
127            zeros += 1;
128            if zeros > 31 {
129                return None;
130            }
131        }
132        let suffix = self.read_bits(zeros)?;
133        Some((1u32 << zeros) - 1 + suffix)
134    }
135
136    pub fn read_se(&mut self) -> Option<i32> {
137        let ue = self.read_ue()?;
138        let v = ((ue + 1) >> 1) as i32;
139        Some(if ue & 1 == 0 { -v } else { v })
140    }
141
142    pub fn byte_align(&mut self) {
143        if self.bit_pos < 8 {
144            self.byte_pos += 1;
145            self.current = self.data.get(self.byte_pos).copied().unwrap_or(0);
146            self.bit_pos = 8;
147        }
148    }
149
150    pub fn remaining_bytes(&self) -> usize {
151        self.data.len().saturating_sub(self.byte_pos)
152    }
153
154    pub fn bits_read(&self) -> usize {
155        self.byte_pos * 8 + (8 - self.bit_pos as usize)
156    }
157
158    /// Remaining bits in the underlying buffer.
159    pub fn bits_left(&self) -> isize {
160        let total = (self.data.len() * 8) as isize;
161        let used = self.bits_read() as isize;
162        total - used
163    }
164
165    pub fn is_empty(&self) -> bool {
166        self.byte_pos >= self.data.len()
167    }
168}