msg_tool\utils/
bit_stream.rs

1//! Bit stream utilities.
2use crate::ext::io::*;
3use anyhow::Result;
4use std::io::{Read, Write};
5
6/// A most significant bit (MSB) bit stream reader.
7pub struct MsbBitStream<T: Read> {
8    /// The input stream to read from.
9    pub m_input: T,
10    m_bits: u32,
11    /// The number of bits currently cached.
12    pub m_cached_bits: u32,
13}
14
15impl<T: Read> MsbBitStream<T> {
16    /// Creates a new MSB bit stream reader.
17    pub fn new(input: T) -> Self {
18        MsbBitStream {
19            m_input: input,
20            m_bits: 0,
21            m_cached_bits: 0,
22        }
23    }
24
25    /// Reads a specified number of bits from the stream.
26    pub fn get_bits(&mut self, count: u32) -> Result<u32> {
27        while self.m_cached_bits < count {
28            let byte = self.m_input.read_u8()?;
29            self.m_bits = (self.m_bits << 8) | byte as u32;
30            self.m_cached_bits += 8;
31        }
32        let mask = (1 << count) - 1;
33        self.m_cached_bits -= count;
34        let result = (self.m_bits >> self.m_cached_bits) & mask;
35        Ok(result)
36    }
37
38    /// Reads the next bit from the stream.
39    pub fn get_next_bit(&mut self) -> Result<bool> {
40        if self.m_cached_bits == 0 {
41            let byte = self.m_input.read_u8()?;
42            self.m_bits = (self.m_bits << 8) | byte as u32;
43            self.m_cached_bits += 8;
44        }
45        self.m_cached_bits -= 1;
46        let bit = (self.m_bits >> self.m_cached_bits) & 1 != 0;
47        Ok(bit)
48    }
49}
50
51/// A most significant bit (MSB) bit writer.
52pub struct MsbBitWriter<'a, T: Write> {
53    /// The output stream to write to.
54    pub writer: &'a mut T,
55    buffer: u32,
56    buffer_size: u32,
57}
58
59impl<'a, T: Write> MsbBitWriter<'a, T> {
60    /// Creates a new MSB bit writer.
61    pub fn new(writer: &'a mut T) -> Self {
62        MsbBitWriter {
63            writer,
64            buffer: 0,
65            buffer_size: 0,
66        }
67    }
68
69    /// Flushes the buffer to the output stream.
70    /// This writes any remaining bits in the buffer to the stream.
71    pub fn flush(&mut self) -> Result<()> {
72        if self.buffer_size > 0 {
73            self.writer
74                .write_u8(((self.buffer << (8 - self.buffer_size)) & 0xFF) as u8)?;
75            self.buffer = 0;
76            self.buffer_size = 0;
77        }
78        Ok(())
79    }
80
81    /// Puts a byte into the bit stream with a specified token width.
82    pub fn put_bits(&mut self, byte: u32, token_width: u8) -> Result<()> {
83        for i in 0..token_width {
84            self.put_bit((byte & (1 << (token_width - 1 - i))) != 0)?;
85        }
86        Ok(())
87    }
88
89    /// Puts a single bit into the bit stream.
90    pub fn put_bit(&mut self, bit: bool) -> Result<()> {
91        self.buffer <<= 1;
92        if bit {
93            self.buffer |= 1;
94        }
95        self.buffer_size += 1;
96        if self.buffer_size == 8 {
97            self.writer.write_u8((self.buffer & 0xFF) as u8)?;
98            self.buffer_size -= 8;
99        }
100        Ok(())
101    }
102}
103
104/// A least significant bit (LSB) bit stream reader.
105pub struct LsbBitStream<T: Read> {
106    /// The input stream to read from.
107    pub m_input: T,
108    m_bits: u32,
109    /// The number of bits currently cached.
110    pub m_cached_bits: u32,
111}
112
113impl<T: Read> LsbBitStream<T> {
114    /// Creates a new LSB bit stream reader.
115    pub fn new(input: T) -> Self {
116        LsbBitStream {
117            m_input: input,
118            m_bits: 0,
119            m_cached_bits: 0,
120        }
121    }
122
123    /// Reads a specified number of bits from the stream.
124    pub fn get_bits(&mut self, mut count: u32) -> Result<u32> {
125        if self.m_cached_bits >= count {
126            let mask = (1 << count) - 1;
127            let value = self.m_bits & mask;
128            self.m_bits >>= count;
129            self.m_cached_bits -= count;
130            Ok(value)
131        } else {
132            let mut value = self.m_bits & ((1 << self.m_cached_bits) - 1);
133            count -= self.m_cached_bits;
134            let mut shift = self.m_cached_bits;
135            self.m_cached_bits = 0;
136            while count >= 8 {
137                let b = self.m_input.read_u8()?;
138                value |= (b as u32) << shift;
139                shift += 8;
140                count -= 8;
141            }
142            if count > 0 {
143                let b = self.m_input.read_u8()?;
144                value |= ((b as u32) & ((1 << count) - 1)) << shift;
145                self.m_bits = b as u32 >> count;
146                self.m_cached_bits = 8 - count;
147            }
148            Ok(value)
149        }
150    }
151
152    /// Reads the next bit from the stream.
153    pub fn get_next_bit(&mut self) -> Result<bool> {
154        Ok(self.get_bits(1)? == 1)
155    }
156}